diff --git a/cake_exports/kmeans/.gitignore b/cake_exports/kmeans/.gitignore new file mode 100644 index 00000000..735d478a --- /dev/null +++ b/cake_exports/kmeans/.gitignore @@ -0,0 +1,8 @@ +__pycache__/ +*.py[cod] +*.egg-info/ +.pytest_cache/ +.ruff_cache/ +results/ +build/ +dist/ diff --git a/cake_exports/kmeans/BENCHMARK_RESULTS.json b/cake_exports/kmeans/BENCHMARK_RESULTS.json new file mode 100644 index 00000000..8558f297 --- /dev/null +++ b/cake_exports/kmeans/BENCHMARK_RESULTS.json @@ -0,0 +1,76379 @@ +{ + "api": "flashlib_cake_kmeans.init(...).compute", + "artifact": "artifacts/generalize_auto_tuning/flash_kmeans_assign_20260620/dispatcher_perf_all_shapes_latest.md", + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_entrypoint": "benchmarks.flash_kmeans_triton_h200_raw_adapter:TritonH20007cfRawAdapter.compute", + "baseline_name": "triton_h200_07cf_raw_adapter_v1", + "baseline_timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e", + "benchmark_registry_baseline_key": "triton_h200_07cf_dual_lane_v1", + "benchmark_registry_baseline_profile": { + "official": { + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_entrypoint": "benchmarks.flash_kmeans_triton_h200_raw_adapter:TritonH20007cfRawAdapter.compute", + "baseline_name": "triton_h200_07cf_raw_adapter_v1", + "baseline_timing_backend_field": "baseline_07cf_adapter_timing_backend", + "baseline_timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e", + "candidate_entrypoint": "flashlib_cake_kmeans.interface:FlashKMeansAssignRuntime.compute", + "candidate_timing_backend_field": "candidate_public_raw_timing_backend", + "candidate_timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion", + "role": "publication", + "speedup_denominator_metric": "candidate_public_raw_synchronized_e2e_ms", + "speedup_metric": "public_raw_e2e_speedup_vs_07cf_adapter", + "speedup_numerator_metric": "baseline_07cf_adapter_synchronized_e2e_ms", + "timing_backend": "cupti" + }, + "parity": { + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_entrypoint": "benchmarks.flash_kmeans_triton_h200:euclid_assign_triton_h200", + "baseline_name": "triton_h200_07cf_precomputed", + "baseline_timing_backend_field": "baseline_07cf_precomputed_timing_backend", + "baseline_timing_boundary": "precomputed_norms_preallocated_output_pinned_07cf_assignment_gpu_span", + "candidate_entrypoint": "flashlib_cake_kmeans.interface:flash_kmeans_assign_prepared", + "candidate_timing_backend_field": "candidate_precomputed_timing_backend", + "candidate_timing_boundary": "precomputed_norms_preallocated_output_prepared_assignment_gpu_span", + "role": "diagnostic_only", + "speedup_denominator_metric": "candidate_precomputed_gpu_span_ms", + "speedup_metric": "precomputed_gpu_speedup_vs_07cf", + "speedup_numerator_metric": "baseline_07cf_precomputed_gpu_span_ms", + "timing_backend": "cupti" + }, + "registry_candidate_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shared_preprocess": { + "implementation_entrypoint": "flashlib_cake_kmeans._row_norm:PreparedBF16PairRowNorm", + "result_source_sha256_field": "preprocess_source_sha256", + "source_sha256": "aa67813cf1cc39b8ae96970a737e926f7b3a65dac63dbeb6362f2dacf066e26e" + } + }, + "benchmark_registry_baseline_sha256": "bdfd30338aa614f09817af1498c1115a0b1732a3340a1acbc172d0e4dd4674c4", + "benchmark_registry_candidate_timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion", + "candidate_timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion", + "evolution_summary": { + "mean_evolution_speedup": 19.980689035087718, + "mean_evolution_tflops": 25.401755701754386, + "min_evolution_speedup": 1.3307, + "row_count": 228, + "unique_shape_count": 225 + }, + "hardware": { + "arch": "sm_100a", + "device": "NVIDIA GB200" + }, + "measured_candidate_entrypoint": "flashlib_cake_kmeans.interface:FlashKMeansAssignRuntime.compute", + "measurement_sessions": [ + { + "adaptive_probe_scope": "separate_nonreportable_cupti_estimation_only", + "alternate_two_pointer_sets": true, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_has_explicit_init": true, + "cache_policy": "synchronize_and_clear_after_each_completed_shape", + "cold_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "id": "a591a388bda74a2890fb7b17b61b2608", + "init_composition": "runtime_init_plus_shared_preprocess_support_each_standalone_lane", + "init_order_policy": "alternate_candidate_baseline_first_by_validation_shard_parity_then_shared_support", + "interleaved": true, + "order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_kernel_parity": { + "baseline_api": "euclid_assign_triton_h200_07cf", + "baseline_norm_policy": "explicit_precomputed_outside_timing", + "baseline_output_policy": "preallocated_outside_timing", + "both_pointer_sets_prepared_before_timing": true, + "candidate_api": "flash_kmeans_assign_prepared", + "candidate_norm_policy": "explicit_precomputed_outside_timing", + "candidate_output_policy": "preallocated_outside_timing" + }, + "public_raw_e2e": { + "assignment_baseline": "frozen_07cf", + "baseline_api": "triton_h200_07cf_raw_adapter_v1.compute(raw_inputs)", + "baseline_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_runtime_initialized_once": true, + "baseline_scratch_policy": "per_shape_per_stream_cached", + "candidate_api": "flashlib_cake_kmeans.init(...).compute(raw_inputs)", + "candidate_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_runtime_initialized_once": true, + "candidate_scratch_policy": "per_shape_per_stream_cached", + "comparison_scope": "complete_raw_input_operators_not_assignment_only", + "fixture_synchronized_before_cold_calls": true, + "preprocess_impl": "flashlib_cake_kmeans._row_norm:PreparedBF16PairRowNorm", + "preprocess_source_sha256": "aa67813cf1cc39b8ae96970a737e926f7b3a65dac63dbeb6362f2dacf066e26e", + "shared_preprocess_cold_compile_attributed_to_lane": null + }, + "randomized_or_interleaved_order": true, + "reportable_timing_collected": true, + "resident_multi_shape_cache_benchmarked": false, + "runtime_instances_reused_across_shapes": true, + "same_cupti_session": true, + "same_process": true, + "scope": "per_shape_single_cupti_activity_session_interleaved_paired_roles", + "sequential_full_sweeps": false, + "timing_blocks": "one_interleaved_cupti_activity_session_per_shape" + }, + { + "adaptive_probe_scope": "separate_nonreportable_cupti_estimation_only", + "alternate_two_pointer_sets": true, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_has_explicit_init": true, + "cache_policy": "synchronize_and_clear_after_each_completed_shape", + "cold_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "id": "41892d8f27794f7aba63e10273932712", + "init_composition": "runtime_init_plus_shared_preprocess_support_each_standalone_lane", + "init_order_policy": "alternate_candidate_baseline_first_by_validation_shard_parity_then_shared_support", + "interleaved": true, + "order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_kernel_parity": { + "baseline_api": "euclid_assign_triton_h200_07cf", + "baseline_norm_policy": "explicit_precomputed_outside_timing", + "baseline_output_policy": "preallocated_outside_timing", + "both_pointer_sets_prepared_before_timing": true, + "candidate_api": "flash_kmeans_assign_prepared", + "candidate_norm_policy": "explicit_precomputed_outside_timing", + "candidate_output_policy": "preallocated_outside_timing" + }, + "public_raw_e2e": { + "assignment_baseline": "frozen_07cf", + "baseline_api": "triton_h200_07cf_raw_adapter_v1.compute(raw_inputs)", + "baseline_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_runtime_initialized_once": true, + "baseline_scratch_policy": "per_shape_per_stream_cached", + "candidate_api": "flashlib_cake_kmeans.init(...).compute(raw_inputs)", + "candidate_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_runtime_initialized_once": true, + "candidate_scratch_policy": "per_shape_per_stream_cached", + "comparison_scope": "complete_raw_input_operators_not_assignment_only", + "fixture_synchronized_before_cold_calls": true, + "preprocess_impl": "flashlib_cake_kmeans._row_norm:PreparedBF16PairRowNorm", + "preprocess_source_sha256": "aa67813cf1cc39b8ae96970a737e926f7b3a65dac63dbeb6362f2dacf066e26e", + "shared_preprocess_cold_compile_attributed_to_lane": null + }, + "randomized_or_interleaved_order": true, + "reportable_timing_collected": true, + "resident_multi_shape_cache_benchmarked": false, + "runtime_instances_reused_across_shapes": true, + "same_cupti_session": true, + "same_process": true, + "scope": "per_shape_single_cupti_activity_session_interleaved_paired_roles", + "sequential_full_sweeps": false, + "timing_blocks": "one_interleaved_cupti_activity_session_per_shape" + }, + { + "adaptive_probe_scope": "separate_nonreportable_cupti_estimation_only", + "alternate_two_pointer_sets": true, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_has_explicit_init": true, + "cache_policy": "synchronize_and_clear_after_each_completed_shape", + "cold_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "init_composition": "runtime_init_plus_shared_preprocess_support_each_standalone_lane", + "init_order_policy": "alternate_candidate_baseline_first_by_validation_shard_parity_then_shared_support", + "interleaved": true, + "order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_kernel_parity": { + "baseline_api": "euclid_assign_triton_h200_07cf", + "baseline_norm_policy": "explicit_precomputed_outside_timing", + "baseline_output_policy": "preallocated_outside_timing", + "both_pointer_sets_prepared_before_timing": true, + "candidate_api": "flash_kmeans_assign_prepared", + "candidate_norm_policy": "explicit_precomputed_outside_timing", + "candidate_output_policy": "preallocated_outside_timing" + }, + "public_raw_e2e": { + "assignment_baseline": "frozen_07cf", + "baseline_api": "triton_h200_07cf_raw_adapter_v1.compute(raw_inputs)", + "baseline_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_runtime_initialized_once": true, + "baseline_scratch_policy": "per_shape_per_stream_cached", + "candidate_api": "flashlib_cake_kmeans.init(...).compute(raw_inputs)", + "candidate_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_runtime_initialized_once": true, + "candidate_scratch_policy": "per_shape_per_stream_cached", + "comparison_scope": "complete_raw_input_operators_not_assignment_only", + "fixture_synchronized_before_cold_calls": true, + "preprocess_impl": "flashlib_cake_kmeans._row_norm:PreparedBF16PairRowNorm", + "preprocess_source_sha256": "aa67813cf1cc39b8ae96970a737e926f7b3a65dac63dbeb6362f2dacf066e26e", + "shared_preprocess_cold_compile_attributed_to_lane": null + }, + "randomized_or_interleaved_order": true, + "reportable_timing_collected": true, + "resident_multi_shape_cache_benchmarked": false, + "runtime_instances_reused_across_shapes": true, + "same_cupti_session": true, + "same_process": true, + "scope": "per_shape_single_cupti_activity_session_interleaved_paired_roles", + "sequential_full_sweeps": false, + "timing_blocks": "one_interleaved_cupti_activity_session_per_shape" + }, + { + "adaptive_probe_scope": "separate_nonreportable_cupti_estimation_only", + "alternate_two_pointer_sets": true, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_has_explicit_init": true, + "cache_policy": "synchronize_and_clear_after_each_completed_shape", + "cold_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "id": "429dbdfd2f734f02819474b9385e8bb8", + "init_composition": "runtime_init_plus_shared_preprocess_support_each_standalone_lane", + "init_order_policy": "alternate_candidate_baseline_first_by_validation_shard_parity_then_shared_support", + "interleaved": true, + "order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_kernel_parity": { + "baseline_api": "euclid_assign_triton_h200_07cf", + "baseline_norm_policy": "explicit_precomputed_outside_timing", + "baseline_output_policy": "preallocated_outside_timing", + "both_pointer_sets_prepared_before_timing": true, + "candidate_api": "flash_kmeans_assign_prepared", + "candidate_norm_policy": "explicit_precomputed_outside_timing", + "candidate_output_policy": "preallocated_outside_timing" + }, + "public_raw_e2e": { + "assignment_baseline": "frozen_07cf", + "baseline_api": "triton_h200_07cf_raw_adapter_v1.compute(raw_inputs)", + "baseline_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_runtime_initialized_once": true, + "baseline_scratch_policy": "per_shape_per_stream_cached", + "candidate_api": "flashlib_cake_kmeans.init(...).compute(raw_inputs)", + "candidate_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_runtime_initialized_once": true, + "candidate_scratch_policy": "per_shape_per_stream_cached", + "comparison_scope": "complete_raw_input_operators_not_assignment_only", + "fixture_synchronized_before_cold_calls": true, + "preprocess_impl": "flashlib_cake_kmeans._row_norm:PreparedBF16PairRowNorm", + "preprocess_source_sha256": "aa67813cf1cc39b8ae96970a737e926f7b3a65dac63dbeb6362f2dacf066e26e", + "shared_preprocess_cold_compile_attributed_to_lane": null + }, + "randomized_or_interleaved_order": true, + "reportable_timing_collected": true, + "resident_multi_shape_cache_benchmarked": false, + "runtime_instances_reused_across_shapes": true, + "same_cupti_session": true, + "same_process": true, + "scope": "per_shape_single_cupti_activity_session_interleaved_paired_roles", + "sequential_full_sweeps": false, + "timing_blocks": "one_interleaved_cupti_activity_session_per_shape" + } + ], + "metadata_only": false, + "official_timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion", + "parity_metric": "precomputed_gpu_speedup_vs_07cf", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_parity_speedup_convention": "precomputed_gpu_speedup_vs_07cf = baseline_07cf_precomputed_gpu_span_ms / candidate_precomputed_gpu_span_ms", + "preprocess_impl": "flashlib_cake_kmeans._row_norm:PreparedBF16PairRowNorm", + "preprocess_source_sha256": "aa67813cf1cc39b8ae96970a737e926f7b3a65dac63dbeb6362f2dacf066e26e", + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "publication_metric": "public_raw_e2e_speedup_vs_07cf_adapter", + "publication_speedup_convention": "public_raw_e2e_speedup_vs_07cf_adapter = baseline_07cf_adapter_synchronized_e2e_ms / candidate_public_raw_synchronized_e2e_ms", + "publication_speedup_metric": "public_raw_e2e_speedup_vs_07cf_adapter", + "publication_timing_backend": "cupti", + "registry_candidate_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "results": [ + { + "B": 1, + "D": 112, + "K": 4096, + "N": 1408, + "baseline_07cf_adapter_bench_iters": 2727, + "baseline_07cf_adapter_gpu_span_ms": 0.088096, + "baseline_07cf_adapter_host_enqueue_ms": 0.139744, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.036928, + "baseline_07cf_adapter_kernel_sum_ms": 0.051137, + "baseline_07cf_adapter_submission_ms": 0.139744, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.184032, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.051137 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.3072, + "synchronized_e2e_ms": 0.342048 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.088096 + }, + "host_enqueue_ms": { + "median": 0.139744 + }, + "inter_kernel_gap_ms": { + "median": 0.036928 + }, + "kernel_sum_ms": { + "median": 0.051137 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2727, + "submission_ms": { + "median": 0.139744 + }, + "synchronized_e2e_ms": { + "median": 0.184032 + } + }, + "baseline_07cf_precomputed_bench_iters": 2574, + "baseline_07cf_precomputed_gpu_span_ms": 0.052353, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042016, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.052353, + "baseline_07cf_precomputed_submission_ms": 0.042016, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.101136, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.052353 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.044064, + "synchronized_e2e_ms": 0.091744 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.052353 + }, + "host_enqueue_ms": { + "median": 0.042016 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.052353 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2574, + "submission_ms": { + "median": 0.042016 + }, + "synchronized_e2e_ms": { + "median": 0.101136 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.6827306935610185, + "submission": 3.3259710586443263, + "synchronized_e2e": 1.8196487897484575 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 356.642449, + "after_init_synchronized_e2e_ms_per_call": 356.705265, + "including_init_host_enqueue_ms_per_call": 390.870356, + "including_init_synchronized_e2e_ms_per_call": 391.050708, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 35.7900145, + "after_init_synchronized_e2e_ms_per_call": 35.8361553, + "including_init_host_enqueue_ms_per_call": 39.212805200000005, + "including_init_synchronized_e2e_ms_per_call": 39.2706996, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 3.7047710499999997, + "after_init_synchronized_e2e_ms_per_call": 3.7492443300000002, + "including_init_host_enqueue_ms_per_call": 4.04705012, + "including_init_synchronized_e2e_ms_per_call": 4.09269876, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.496246705, + "after_init_synchronized_e2e_ms_per_call": 0.540553233, + "including_init_host_enqueue_ms_per_call": 0.530474612, + "including_init_synchronized_e2e_ms_per_call": 0.574898676, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 356.642449, + "synchronized_e2e_ms": 356.705265, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.529889, + "median": 0.088096, + "min": 0.081024, + "p90": 0.107456, + "sample_count": 2727 + }, + "host_enqueue_ms": { + "max": 1.376257, + "median": 0.139744, + "min": 0.118721, + "p90": 0.2122496, + "sample_count": 2727 + }, + "sample_count": 2727, + "synchronized_e2e_ms": { + "max": 1.652769, + "median": 0.184032, + "min": 0.164257, + "p90": 0.2479424, + "sample_count": 2727 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 2574, + "candidate_precomputed_gpu_span_ms": 0.039072, + "candidate_precomputed_host_enqueue_ms": 0.052816, + "candidate_precomputed_inter_kernel_gap_ms": 0.004768, + "candidate_precomputed_kernel_sum_ms": 0.034208, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.052816, + "candidate_precomputed_synchronized_e2e_ms": 0.073504, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.034208 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.043392, + "synchronized_e2e_ms": 0.064672 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.039072 + }, + "host_enqueue_ms": { + "median": 0.052816 + }, + "inter_kernel_gap_ms": { + "median": 0.004768 + }, + "kernel_sum_ms": { + "median": 0.034208 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2574, + "submission_ms": { + "median": 0.052816 + }, + "synchronized_e2e_ms": { + "median": 0.073504 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295c3ef0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282af410" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 0.941031941031941, + "submission": 0.807028173280824, + "synchronized_e2e": 1.1144971702220288 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2727, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.036768, + "candidate_public_raw_host_enqueue_ms": 0.042624, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.03696, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.042624, + "candidate_public_raw_synchronized_e2e_ms": 0.08192, + "candidate_public_raw_tflops_from_gpu_span": 35.135053089643165, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.036768 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.067936, + "synchronized_e2e_ms": 0.09168 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.036768 + }, + "host_enqueue_ms": { + "median": 0.042624 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.03696 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2727, + "submission_ms": { + "median": 0.042624 + }, + "synchronized_e2e_ms": { + "median": 0.08192 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 79.679571, + "after_init_synchronized_e2e_ms_per_call": 79.714707, + "including_init_host_enqueue_ms_per_call": 114.257111, + "including_init_synchronized_e2e_ms_per_call": 527.802499, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 8.0063187, + "after_init_synchronized_e2e_ms_per_call": 8.0451987, + "including_init_host_enqueue_ms_per_call": 11.4640727, + "including_init_synchronized_e2e_ms_per_call": 52.853977900000004, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8389934699999999, + "after_init_synchronized_e2e_ms_per_call": 0.87824787, + "including_init_host_enqueue_ms_per_call": 1.1847688699999999, + "including_init_synchronized_e2e_ms_per_call": 5.35912579, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.12226094699999998, + "after_init_synchronized_e2e_ms_per_call": 0.16155278700000003, + "including_init_host_enqueue_ms_per_call": 0.156838487, + "including_init_synchronized_e2e_ms_per_call": 0.609640579, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 79.679571, + "synchronized_e2e_ms": 79.714707, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.03744, + "median": 0.036768, + "min": 0.036384, + "p90": 0.036928, + "sample_count": 2727 + }, + "host_enqueue_ms": { + "max": 24.282201, + "median": 0.042624, + "min": 0.031232, + "p90": 0.06411520000000001, + "sample_count": 2727 + }, + "sample_count": 2727, + "synchronized_e2e_ms": { + "max": 44.927375, + "median": 0.08192, + "min": 0.071712, + "p90": 0.10037760000000001, + "sample_count": 2727 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.069984, + "submission_ms": 0.069984, + "synchronized_e2e_ms": 0.117088 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.044064, + "submission_ms": 0.044064, + "synchronized_e2e_ms": 0.091744 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.3072, + "submission_ms": 0.3072, + "synchronized_e2e_ms": 0.342048 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 356.642449, + "submission_ms": 356.642449, + "synchronized_e2e_ms": 356.705265 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.08064, + "submission_ms": 0.08064, + "synchronized_e2e_ms": 0.095392 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.043392, + "submission_ms": 0.043392, + "synchronized_e2e_ms": 0.064672 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.205058, + "submission_ms": 1.205058, + "synchronized_e2e_ms": 1.228162 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.330177, + "submission_ms": 1.330177, + "synchronized_e2e_ms": 1.353153 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.067936, + "submission_ms": 0.067936, + "synchronized_e2e_ms": 0.09168 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 79.679571, + "submission_ms": 79.679571, + "synchronized_e2e_ms": 79.714707 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.866559, + "evolution_kernel_ms": 0.196192, + "evolution_speedup": 4.4169, + "evolution_tflops": 6.5846, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_1d49_d112_highk_tail_b1_n1408_k4096_d112", + "measurement_order": [ + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 10602, + "measurement_schedule_sha256": "e41f6a73be7affd34f7b9e49d52d3bf88463da79f49e2a3a2a041be117cab67c", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 2574, + "public_pair_count": 2727, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 2574, + "baseline_public_raw": 2727, + "candidate_precomputed": 2574, + "candidate_public_raw": 2727 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2122 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:adjacent_1d49_d112_highk_tail_b1_n1408_k4096_d112", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.3399109336609334, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.2464843749999996, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 4.4747735822449926, + "including_init_synchronized_e2e_speedup": 0.7409034795039877, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 4.454353041647064, + "including_init_synchronized_e2e_speedup": 0.7430036708741273, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 4.269004751471814, + "including_init_synchronized_e2e_speedup": 0.7636877581110109, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 3.345985191824638, + "including_init_synchronized_e2e_speedup": 0.9430124827697862, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.395996518711923, + "hot_synchronized_e2e_speedup": 2.2464843749999996, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 1491121, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_1d49_d112_highk_tail_b1_n1408_k4096_d112", + "source": "tail_divisibility", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 128, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 4, + "num_warps": 4 + } + }, + { + "B": 5, + "D": 128, + "K": 512, + "N": 7296, + "baseline_07cf_adapter_bench_iters": 2730, + "baseline_07cf_adapter_gpu_span_ms": 0.06488, + "baseline_07cf_adapter_host_enqueue_ms": 0.1553445, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.023328, + "baseline_07cf_adapter_kernel_sum_ms": 0.041504, + "baseline_07cf_adapter_submission_ms": 0.1553445, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.174192, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.041504 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.297537, + "synchronized_e2e_ms": 0.317185 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.06488 + }, + "host_enqueue_ms": { + "median": 0.1553445 + }, + "inter_kernel_gap_ms": { + "median": 0.023328 + }, + "kernel_sum_ms": { + "median": 0.041504 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2730, + "submission_ms": { + "median": 0.1553445 + }, + "synchronized_e2e_ms": { + "median": 0.174192 + } + }, + "baseline_07cf_precomputed_bench_iters": 7371, + "baseline_07cf_precomputed_gpu_span_ms": 0.018176, + "baseline_07cf_precomputed_host_enqueue_ms": 0.04336, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.018176, + "baseline_07cf_precomputed_submission_ms": 0.04336, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.066976, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.018176 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.048768, + "synchronized_e2e_ms": 0.070496 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.018176 + }, + "host_enqueue_ms": { + "median": 0.04336 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.018176 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 7371, + "submission_ms": { + "median": 0.04336 + }, + "synchronized_e2e_ms": { + "median": 0.066976 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 3.569542253521126, + "submission": 3.582668357933579, + "synchronized_e2e": 2.6008122312470143 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 349.246249, + "after_init_synchronized_e2e_ms_per_call": 349.282089, + "including_init_host_enqueue_ms_per_call": 383.695693, + "including_init_synchronized_e2e_ms_per_call": 833.323583, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 35.06443495, + "after_init_synchronized_e2e_ms_per_call": 35.0849817, + "including_init_host_enqueue_ms_per_call": 38.50937935, + "including_init_synchronized_e2e_ms_per_call": 83.4891311, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 3.6462535449999995, + "after_init_synchronized_e2e_ms_per_call": 3.66527097, + "including_init_host_enqueue_ms_per_call": 3.990747985, + "including_init_synchronized_e2e_ms_per_call": 8.50568591, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.5044354045, + "after_init_synchronized_e2e_ms_per_call": 0.523299897, + "including_init_host_enqueue_ms_per_call": 0.5388848485, + "including_init_synchronized_e2e_ms_per_call": 1.007341391, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 349.246249, + "synchronized_e2e_ms": 349.282089, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 42.738987, + "median": 0.06488, + "min": 0.054272, + "p90": 0.086784, + "sample_count": 2730 + }, + "host_enqueue_ms": { + "max": 85.932505, + "median": 0.1553445, + "min": 0.126528, + "p90": 0.20838079999999998, + "sample_count": 2730 + }, + "sample_count": 2730, + "synchronized_e2e_ms": { + "max": 85.974745, + "median": 0.174192, + "min": 0.143072, + "p90": 0.23240639999999999, + "sample_count": 2730 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 7371, + "candidate_precomputed_gpu_span_ms": 0.013728, + "candidate_precomputed_host_enqueue_ms": 0.040448, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.013728, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.040448, + "candidate_precomputed_synchronized_e2e_ms": 0.051808, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.013728 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.035936, + "synchronized_e2e_ms": 0.047232 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.013728 + }, + "host_enqueue_ms": { + "median": 0.040448 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.013728 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 7371, + "submission_ms": { + "median": 0.040448 + }, + "synchronized_e2e_ms": { + "median": 0.051808 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295bfe00", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282afe00" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.6829836829836826, + "submission": 1.2120253164556962, + "synchronized_e2e": 1.650401482396541 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 2730, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.036832, + "candidate_public_raw_host_enqueue_ms": 0.049024, + "candidate_public_raw_inter_kernel_gap_ms": 0.00016, + "candidate_public_raw_kernel_sum_ms": 0.03664, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.049024, + "candidate_public_raw_synchronized_e2e_ms": 0.085504, + "candidate_public_raw_tflops_from_gpu_span": 129.8193570807993, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.03664 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.49648, + "synchronized_e2e_ms": 0.519168 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.036832 + }, + "host_enqueue_ms": { + "median": 0.049024 + }, + "inter_kernel_gap_ms": { + "median": 0.00016 + }, + "kernel_sum_ms": { + "median": 0.03664 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2730, + "submission_ms": { + "median": 0.049024 + }, + "synchronized_e2e_ms": { + "median": 0.085504 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 39.483945, + "after_init_synchronized_e2e_ms_per_call": 39.514729, + "including_init_host_enqueue_ms_per_call": 74.35802899999999, + "including_init_synchronized_e2e_ms_per_call": 74.45780500000001, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 3.9925160999999996, + "after_init_synchronized_e2e_ms_per_call": 4.0284265, + "including_init_host_enqueue_ms_per_call": 7.479924499999998, + "including_init_synchronized_e2e_ms_per_call": 7.522734100000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.44337320999999996, + "after_init_synchronized_e2e_ms_per_call": 0.47979625, + "including_init_host_enqueue_ms_per_call": 0.7921140499999999, + "including_init_synchronized_e2e_ms_per_call": 0.82922701, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.08845892100000001, + "after_init_synchronized_e2e_ms_per_call": 0.124933225, + "including_init_host_enqueue_ms_per_call": 0.12333300499999998, + "including_init_synchronized_e2e_ms_per_call": 0.159876301, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 39.483945, + "synchronized_e2e_ms": 39.514729, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.0376, + "median": 0.036832, + "min": 0.036288, + "p90": 0.03712, + "sample_count": 2730 + }, + "host_enqueue_ms": { + "max": 17.227922, + "median": 0.049024, + "min": 0.037664, + "p90": 0.07346239999999998, + "sample_count": 2730 + }, + "sample_count": 2730, + "synchronized_e2e_ms": { + "max": 20.491349, + "median": 0.085504, + "min": 0.075808, + "p90": 0.10778559999999998, + "sample_count": 2730 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.084865, + "submission_ms": 0.084865, + "synchronized_e2e_ms": 0.105505 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.048768, + "submission_ms": 0.048768, + "synchronized_e2e_ms": 0.070496 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.297537, + "submission_ms": 0.297537, + "synchronized_e2e_ms": 0.317185 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 349.246249, + "submission_ms": 349.246249, + "synchronized_e2e_ms": 349.282089 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.061088, + "submission_ms": 0.061088, + "synchronized_e2e_ms": 0.074624 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.035936, + "submission_ms": 0.035936, + "synchronized_e2e_ms": 0.047232 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.141345, + "submission_ms": 1.141345, + "synchronized_e2e_ms": 1.160833 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.001345, + "submission_ms": 1.001345, + "synchronized_e2e_ms": 1.018945 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.49648, + "submission_ms": 0.49648, + "synchronized_e2e_ms": 0.519168 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 39.483945, + "submission_ms": 39.483945, + "synchronized_e2e_ms": 39.514729 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.339855, + "evolution_kernel_ms": 0.152639, + "evolution_speedup": 2.2265, + "evolution_tflops": 31.3255, + "expected_route": "aligned_weave_v10_fallback", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_1d49_d128_forced_fallback_b5_n7296_k512_d128", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 20202, + "measurement_schedule_sha256": "3606b84b5a67854996dd26875b960e2a3b042030c864de63d4ab6eb8ff67de46", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 7371, + "public_pair_count": 2730, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 7371, + "baseline_public_raw": 2730, + "candidate_precomputed": 7371, + "candidate_public_raw": 2730 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 4042 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:adjacent_1d49_d128_forced_fallback_b5_n7296_k512_d128", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.324009324009324, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.037238023952096, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 8.839288484048568, + "including_init_synchronized_e2e_speedup": 11.191890265902948, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 8.709351331096645, + "including_init_synchronized_e2e_speedup": 11.098243004494867, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 7.63922387888609, + "including_init_synchronized_e2e_speedup": 10.257367171385313, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 4.18863674574958, + "including_init_synchronized_e2e_speedup": 6.300754925522076, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.761511728931364, + "hot_synchronized_e2e_speedup": 2.037238023952096, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 1491281, + "selected_route": "aligned_weave_v10_fallback", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_1d49_d128_forced_fallback_b5_n7296_k512_d128", + "source": "forced_fallback", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 5, + "D": 224, + "K": 768, + "N": 5632, + "baseline_07cf_adapter_bench_iters": 1573, + "baseline_07cf_adapter_gpu_span_ms": 0.072224, + "baseline_07cf_adapter_host_enqueue_ms": 0.143968, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.0056, + "baseline_07cf_adapter_kernel_sum_ms": 0.066496, + "baseline_07cf_adapter_submission_ms": 0.143968, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.168801, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.066496 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.247425, + "synchronized_e2e_ms": 0.267873 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.072224 + }, + "host_enqueue_ms": { + "median": 0.143968 + }, + "inter_kernel_gap_ms": { + "median": 0.0056 + }, + "kernel_sum_ms": { + "median": 0.066496 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1573, + "submission_ms": { + "median": 0.143968 + }, + "synchronized_e2e_ms": { + "median": 0.168801 + } + }, + "baseline_07cf_precomputed_bench_iters": 3450, + "baseline_07cf_precomputed_gpu_span_ms": 0.031008, + "baseline_07cf_precomputed_host_enqueue_ms": 0.0408805, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.031008, + "baseline_07cf_precomputed_submission_ms": 0.0408805, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.0774405, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.031008 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.042688, + "synchronized_e2e_ms": 0.088288 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.031008 + }, + "host_enqueue_ms": { + "median": 0.0408805 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.031008 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3450, + "submission_ms": { + "median": 0.0408805 + }, + "synchronized_e2e_ms": { + "median": 0.0774405 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.329205366357069, + "submission": 3.521679040129157, + "synchronized_e2e": 2.1797509055339264 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 364.953146, + "after_init_synchronized_e2e_ms_per_call": 365.01561, + "including_init_host_enqueue_ms_per_call": 400.504414, + "including_init_synchronized_e2e_ms_per_call": 400.646847, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 36.6248858, + "after_init_synchronized_e2e_ms_per_call": 36.653481899999996, + "including_init_host_enqueue_ms_per_call": 40.1800126, + "including_init_synchronized_e2e_ms_per_call": 40.216605599999994, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 3.79205978, + "after_init_synchronized_e2e_ms_per_call": 3.81726909, + "including_init_host_enqueue_ms_per_call": 4.14757246, + "including_init_synchronized_e2e_ms_per_call": 4.173581459999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.508777178, + "after_init_synchronized_e2e_ms_per_call": 0.5336478090000001, + "including_init_host_enqueue_ms_per_call": 0.544328446, + "including_init_synchronized_e2e_ms_per_call": 0.569279046, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 364.953146, + "synchronized_e2e_ms": 365.01561, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 45.355404, + "median": 0.072224, + "min": 0.066656, + "p90": 0.0916672, + "sample_count": 1573 + }, + "host_enqueue_ms": { + "max": 45.726192, + "median": 0.143968, + "min": 0.126304, + "p90": 0.19106559999999997, + "sample_count": 1573 + }, + "sample_count": 1573, + "synchronized_e2e_ms": { + "max": 45.857936, + "median": 0.168801, + "min": 0.152416, + "p90": 0.21676819999999994, + "sample_count": 1573 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3450, + "candidate_precomputed_gpu_span_ms": 0.028832, + "candidate_precomputed_host_enqueue_ms": 0.050752, + "candidate_precomputed_inter_kernel_gap_ms": 0.001952, + "candidate_precomputed_kernel_sum_ms": 0.02688, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.050752, + "candidate_precomputed_synchronized_e2e_ms": 0.06192, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.02688 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.046528, + "synchronized_e2e_ms": 0.058464 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.028832 + }, + "host_enqueue_ms": { + "median": 0.050752 + }, + "inter_kernel_gap_ms": { + "median": 0.001952 + }, + "kernel_sum_ms": { + "median": 0.02688 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3450, + "submission_ms": { + "median": 0.050752 + }, + "synchronized_e2e_ms": { + "median": 0.06192 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295c33e0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282af5c0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.205327413984462, + "submission": 0.8480453972257251, + "synchronized_e2e": 1.7436854005167957 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 1573, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.063584, + "candidate_public_raw_host_enqueue_ms": 0.04304, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.063808, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.04304, + "candidate_public_raw_synchronized_e2e_ms": 0.107969, + "candidate_public_raw_tflops_from_gpu_span": 152.3786210367388, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.063584 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.061376, + "synchronized_e2e_ms": 0.106048 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.063584 + }, + "host_enqueue_ms": { + "median": 0.04304 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.063808 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1573, + "submission_ms": { + "median": 0.04304 + }, + "synchronized_e2e_ms": { + "median": 0.107969 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 80.469491, + "after_init_synchronized_e2e_ms_per_call": 80.500755, + "including_init_host_enqueue_ms_per_call": 116.307511, + "including_init_synchronized_e2e_ms_per_call": 564.600968, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 8.085685100000001, + "after_init_synchronized_e2e_ms_per_call": 8.1472476, + "including_init_host_enqueue_ms_per_call": 11.669487100000001, + "including_init_synchronized_e2e_ms_per_call": 56.5572689, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8473045100000001, + "after_init_synchronized_e2e_ms_per_call": 0.9118968599999999, + "including_init_host_enqueue_ms_per_call": 1.20568471, + "including_init_synchronized_e2e_ms_per_call": 5.75289899, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.123466451, + "after_init_synchronized_e2e_ms_per_call": 0.188361786, + "including_init_host_enqueue_ms_per_call": 0.159304471, + "including_init_synchronized_e2e_ms_per_call": 0.672461999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 80.469491, + "synchronized_e2e_ms": 80.500755, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.06448, + "median": 0.063584, + "min": 0.062944, + "p90": 0.063904, + "sample_count": 1573 + }, + "host_enqueue_ms": { + "max": 0.41008, + "median": 0.04304, + "min": 0.034176, + "p90": 0.06255359999999999, + "sample_count": 1573 + }, + "sample_count": 1573, + "synchronized_e2e_ms": { + "max": 0.493248, + "median": 0.107969, + "min": 0.101344, + "p90": 0.1249344, + "sample_count": 1573 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.068256, + "submission_ms": 0.068256, + "synchronized_e2e_ms": 0.098784 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.042688, + "submission_ms": 0.042688, + "synchronized_e2e_ms": 0.088288 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.247425, + "submission_ms": 0.247425, + "synchronized_e2e_ms": 0.267873 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 364.953146, + "submission_ms": 364.953146, + "synchronized_e2e_ms": 365.01561 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.077601, + "submission_ms": 0.077601, + "synchronized_e2e_ms": 0.093057 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.046528, + "submission_ms": 0.046528, + "synchronized_e2e_ms": 0.058464 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.240193, + "submission_ms": 1.240193, + "synchronized_e2e_ms": 1.262977 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.271873, + "submission_ms": 1.271873, + "synchronized_e2e_ms": 1.293377 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.061376, + "submission_ms": 0.061376, + "synchronized_e2e_ms": 0.106048 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 80.469491, + "submission_ms": 80.469491, + "synchronized_e2e_ms": 80.500755 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.465423, + "evolution_kernel_ms": 0.186336, + "evolution_speedup": 2.4978, + "evolution_tflops": 51.9966, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_1d49_d224_highn_overlap_b5_n5632_k768_d224", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 10046, + "measurement_schedule_sha256": "cde7ef8285260e28350c6ee33ce717692c70460cf386d71fad05a3cb1b1498d6", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3450, + "public_pair_count": 1573, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3450, + "baseline_public_raw": 1573, + "candidate_precomputed": 3450, + "candidate_public_raw": 1573 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2010 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:adjacent_1d49_d224_highn_overlap_b5_n5632_k768_d224", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.0754716981132075, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.5634209819485223, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 4.534312876941341, + "including_init_synchronized_e2e_speedup": 0.7096106271642099, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 4.498879093842685, + "including_init_synchronized_e2e_speedup": 0.7110775746811211, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 4.186075484457749, + "including_init_synchronized_e2e_speedup": 0.725474489862371, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.8331001756375365, + "including_init_synchronized_e2e_speedup": 0.8465594291522189, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.135883241066935, + "hot_synchronized_e2e_speedup": 1.5634209819485223, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 1492241, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_1d49_d224_highn_overlap_b5_n5632_k768_d224", + "source": "guard_overlap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 3, + "D": 288, + "K": 2048, + "N": 1152, + "baseline_07cf_adapter_bench_iters": 1327, + "baseline_07cf_adapter_gpu_span_ms": 0.100448, + "baseline_07cf_adapter_host_enqueue_ms": 0.1528, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.032672, + "baseline_07cf_adapter_kernel_sum_ms": 0.067776, + "baseline_07cf_adapter_submission_ms": 0.1528, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.203296, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.067776 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.387297, + "synchronized_e2e_ms": 0.431105 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.100448 + }, + "host_enqueue_ms": { + "median": 0.1528 + }, + "inter_kernel_gap_ms": { + "median": 0.032672 + }, + "kernel_sum_ms": { + "median": 0.067776 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1327, + "submission_ms": { + "median": 0.1528 + }, + "synchronized_e2e_ms": { + "median": 0.203296 + } + }, + "baseline_07cf_precomputed_bench_iters": 1564, + "baseline_07cf_precomputed_gpu_span_ms": 0.06384, + "baseline_07cf_precomputed_host_enqueue_ms": 0.046592, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.06384, + "baseline_07cf_precomputed_submission_ms": 0.046592, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.117184, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.06384 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.045184, + "synchronized_e2e_ms": 0.100768 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.06384 + }, + "host_enqueue_ms": { + "median": 0.046592 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.06384 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1564, + "submission_ms": { + "median": 0.046592 + }, + "synchronized_e2e_ms": { + "median": 0.117184 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.5734335839598999, + "submission": 3.2795329670329667, + "synchronized_e2e": 1.7348443473511743 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 340.503617, + "after_init_synchronized_e2e_ms_per_call": 340.593569, + "including_init_host_enqueue_ms_per_call": 377.438055, + "including_init_synchronized_e2e_ms_per_call": 795.704888, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 34.187881700000005, + "after_init_synchronized_e2e_ms_per_call": 34.242323299999995, + "including_init_host_enqueue_ms_per_call": 37.8813255, + "including_init_synchronized_e2e_ms_per_call": 79.75345519999999, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 3.5563081700000003, + "after_init_synchronized_e2e_ms_per_call": 3.60719873, + "including_init_host_enqueue_ms_per_call": 3.9256525500000006, + "including_init_synchronized_e2e_ms_per_call": 8.15831192, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.493150817, + "after_init_synchronized_e2e_ms_per_call": 0.543686273, + "including_init_host_enqueue_ms_per_call": 0.5300852549999999, + "including_init_synchronized_e2e_ms_per_call": 0.998797592, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 340.503617, + "synchronized_e2e_ms": 340.593569, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 72.71009, + "median": 0.100448, + "min": 0.090656, + "p90": 0.1122688, + "sample_count": 1327 + }, + "host_enqueue_ms": { + "max": 72.808939, + "median": 0.1528, + "min": 0.128096, + "p90": 0.181735, + "sample_count": 1327 + }, + "sample_count": 1327, + "synchronized_e2e_ms": { + "max": 72.887915, + "median": 0.203296, + "min": 0.180096, + "p90": 0.2293444, + "sample_count": 1327 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 1564, + "candidate_precomputed_gpu_span_ms": 0.077152, + "candidate_precomputed_host_enqueue_ms": 0.0426245, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.077152, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.0426245, + "candidate_precomputed_synchronized_e2e_ms": 0.113872, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.077152 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.038016, + "synchronized_e2e_ms": 0.082272 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.077152 + }, + "host_enqueue_ms": { + "median": 0.0426245 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.077152 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1564, + "submission_ms": { + "median": 0.0426245 + }, + "synchronized_e2e_ms": { + "median": 0.113872 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282adfd0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295c3d10" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 0.9772008502695977, + "submission": 1.1193562387828595, + "synchronized_e2e": 1.0909090909090908 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 1327, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.075393, + "candidate_public_raw_host_enqueue_ms": 0.047712, + "candidate_public_raw_inter_kernel_gap_ms": 0.000192, + "candidate_public_raw_kernel_sum_ms": 0.0752, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.047712, + "candidate_public_raw_synchronized_e2e_ms": 0.124224, + "candidate_public_raw_tflops_from_gpu_span": 54.07482774262862, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.0752 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.623072, + "synchronized_e2e_ms": 0.660192 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.075393 + }, + "host_enqueue_ms": { + "median": 0.047712 + }, + "inter_kernel_gap_ms": { + "median": 0.000192 + }, + "kernel_sum_ms": { + "median": 0.0752 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1327, + "submission_ms": { + "median": 0.047712 + }, + "synchronized_e2e_ms": { + "median": 0.124224 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 44.543214, + "after_init_synchronized_e2e_ms_per_call": 44.577806, + "including_init_host_enqueue_ms_per_call": 81.901237, + "including_init_synchronized_e2e_ms_per_call": 82.013429, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 4.4972622, + "after_init_synchronized_e2e_ms_per_call": 4.5695822, + "including_init_host_enqueue_ms_per_call": 8.2330645, + "including_init_synchronized_e2e_ms_per_call": 8.3131445, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.49266701999999996, + "after_init_synchronized_e2e_ms_per_call": 0.56875982, + "including_init_host_enqueue_ms_per_call": 0.86624725, + "including_init_synchronized_e2e_ms_per_call": 0.94311605, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.09220750200000001, + "after_init_synchronized_e2e_ms_per_call": 0.168677582, + "including_init_host_enqueue_ms_per_call": 0.129565525, + "including_init_synchronized_e2e_ms_per_call": 0.206113205, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 44.543214, + "synchronized_e2e_ms": 44.577806, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.076352, + "median": 0.075393, + "min": 0.074752, + "p90": 0.075776, + "sample_count": 1327 + }, + "host_enqueue_ms": { + "max": 0.150944, + "median": 0.047712, + "min": 0.038528, + "p90": 0.05972480000000001, + "sample_count": 1327 + }, + "sample_count": 1327, + "synchronized_e2e_ms": { + "max": 0.205888, + "median": 0.124224, + "min": 0.115136, + "p90": 0.13310760000000002, + "sample_count": 1327 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.069664, + "submission_ms": 0.069664, + "synchronized_e2e_ms": 0.123232 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.045184, + "submission_ms": 0.045184, + "synchronized_e2e_ms": 0.100768 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.387297, + "submission_ms": 0.387297, + "synchronized_e2e_ms": 0.431105 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 340.503617, + "submission_ms": 340.503617, + "synchronized_e2e_ms": 340.593569 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.073472, + "submission_ms": 0.073472, + "synchronized_e2e_ms": 0.104096 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.038016, + "submission_ms": 0.038016, + "synchronized_e2e_ms": 0.082272 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.355906, + "submission_ms": 1.355906, + "synchronized_e2e_ms": 1.37565 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.030721, + "submission_ms": 1.030721, + "synchronized_e2e_ms": 1.049697 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.623072, + "submission_ms": 0.623072, + "synchronized_e2e_ms": 0.660192 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 44.543214, + "submission_ms": 44.543214, + "synchronized_e2e_ms": 44.577806 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 1.007006, + "evolution_kernel_ms": 0.207744, + "evolution_speedup": 4.8473, + "evolution_tflops": 19.6245, + "expected_route": "d288_parent_splitk_hybrid_20260629_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_1d49_d288_low_n_heldout_b3_n1152_k2048_d288", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 5782, + "measurement_schedule_sha256": "c7dea955c8d7629485113fca414ea9e6c476f9bf4be85e55d1c9fb27ad7fc619", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 1564, + "public_pair_count": 1327, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 1564, + "baseline_public_raw": 1327, + "candidate_precomputed": 1564, + "candidate_public_raw": 1327 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 1158 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:adjacent_1d49_d288_low_n_heldout_b3_n1152_k2048_d288", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.8274574865201161, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.6365275631117981, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 7.640429163337468, + "including_init_synchronized_e2e_speedup": 9.702129245199588, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 7.493534813751681, + "including_init_synchronized_e2e_speedup": 9.593656792565074, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 6.342217933045974, + "including_init_synchronized_e2e_speedup": 8.65037968551166, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 3.2232278086604307, + "including_init_synchronized_e2e_speedup": 4.8458690067916805, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.3323252821880014, + "hot_synchronized_e2e_speedup": 1.6365275631117981, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 1492881, + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_1d49_d288_low_n_heldout_b3_n1152_k2048_d288", + "source": "heldout_neighborhood", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 352, + "K": 8192, + "N": 1024, + "baseline_07cf_adapter_bench_iters": 941, + "baseline_07cf_adapter_gpu_span_ms": 0.248736, + "baseline_07cf_adapter_host_enqueue_ms": 0.14032, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.013152, + "baseline_07cf_adapter_kernel_sum_ms": 0.235521, + "baseline_07cf_adapter_submission_ms": 0.14032, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.345057, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.235521 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.232288, + "synchronized_e2e_ms": 0.419649 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.248736 + }, + "host_enqueue_ms": { + "median": 0.14032 + }, + "inter_kernel_gap_ms": { + "median": 0.013152 + }, + "kernel_sum_ms": { + "median": 0.235521 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 941, + "submission_ms": { + "median": 0.14032 + }, + "synchronized_e2e_ms": { + "median": 0.345057 + } + }, + "baseline_07cf_precomputed_bench_iters": 1193, + "baseline_07cf_precomputed_gpu_span_ms": 0.246304, + "baseline_07cf_precomputed_host_enqueue_ms": 0.041376, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.246304, + "baseline_07cf_precomputed_submission_ms": 0.041376, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.294656, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.246304 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.043264, + "synchronized_e2e_ms": 0.233312 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.246304 + }, + "host_enqueue_ms": { + "median": 0.041376 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.246304 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1193, + "submission_ms": { + "median": 0.041376 + }, + "synchronized_e2e_ms": { + "median": 0.294656 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.009873976874107, + "submission": 3.391337973704563, + "synchronized_e2e": 1.1710503095134666 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 6.955047, + "after_init_synchronized_e2e_ms_per_call": 7.135847, + "including_init_host_enqueue_ms_per_call": 41.182954, + "including_init_synchronized_e2e_ms_per_call": 41.48129, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8217926999999999, + "after_init_synchronized_e2e_ms_per_call": 1.024136, + "including_init_host_enqueue_ms_per_call": 4.244583400000001, + "including_init_synchronized_e2e_ms_per_call": 4.4586803, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.20846727, + "after_init_synchronized_e2e_ms_per_call": 0.41296489999999997, + "including_init_host_enqueue_ms_per_call": 0.55074634, + "including_init_synchronized_e2e_ms_per_call": 0.75641933, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.147134727, + "after_init_synchronized_e2e_ms_per_call": 0.35184779000000005, + "including_init_host_enqueue_ms_per_call": 0.181362634, + "including_init_synchronized_e2e_ms_per_call": 0.386193233, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 6.955047, + "synchronized_e2e_ms": 7.135847, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.477312, + "median": 0.248736, + "min": 0.242497, + "p90": 0.261761, + "sample_count": 941 + }, + "host_enqueue_ms": { + "max": 0.445313, + "median": 0.14032, + "min": 0.125792, + "p90": 0.184833, + "sample_count": 941 + }, + "sample_count": 941, + "synchronized_e2e_ms": { + "max": 0.622529, + "median": 0.345057, + "min": 0.331841, + "p90": 0.385121, + "sample_count": 941 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 1193, + "candidate_precomputed_gpu_span_ms": 0.083584, + "candidate_precomputed_host_enqueue_ms": 0.04784, + "candidate_precomputed_inter_kernel_gap_ms": 0.00208, + "candidate_precomputed_kernel_sum_ms": 0.081504, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.04784, + "candidate_precomputed_synchronized_e2e_ms": 0.11712, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.081504 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.04848, + "synchronized_e2e_ms": 0.099712 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.083584 + }, + "host_enqueue_ms": { + "median": 0.04784 + }, + "inter_kernel_gap_ms": { + "median": 0.00208 + }, + "kernel_sum_ms": { + "median": 0.081504 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1193, + "submission_ms": { + "median": 0.04784 + }, + "synchronized_e2e_ms": { + "median": 0.11712 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04734620", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04734b60" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.2699081163859112, + "submission": 0.9438127090301003, + "synchronized_e2e": 1.313934426229508 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 941, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.106144, + "candidate_public_raw_host_enqueue_ms": 0.045152, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.106336, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.045152, + "candidate_public_raw_synchronized_e2e_ms": 0.153888, + "candidate_public_raw_tflops_from_gpu_span": 111.27487247512813, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.106144 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.620353, + "synchronized_e2e_ms": 0.705057 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.106144 + }, + "host_enqueue_ms": { + "median": 0.045152 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.106336 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 941, + "submission_ms": { + "median": 0.045152 + }, + "synchronized_e2e_ms": { + "median": 0.153888 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 75.892303, + "after_init_synchronized_e2e_ms_per_call": 75.966031, + "including_init_host_enqueue_ms_per_call": 110.469843, + "including_init_synchronized_e2e_ms_per_call": 524.053823, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 7.6298671, + "after_init_synchronized_e2e_ms_per_call": 7.735102299999999, + "including_init_host_enqueue_ms_per_call": 11.0876211, + "including_init_synchronized_e2e_ms_per_call": 52.5438815, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.80362351, + "after_init_synchronized_e2e_ms_per_call": 0.9120094299999999, + "including_init_host_enqueue_ms_per_call": 1.14939891, + "including_init_synchronized_e2e_ms_per_call": 5.39288735, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.120999151, + "after_init_synchronized_e2e_ms_per_call": 0.22970014300000002, + "including_init_host_enqueue_ms_per_call": 0.155576691, + "including_init_synchronized_e2e_ms_per_call": 0.6777879349999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 75.892303, + "synchronized_e2e_ms": 75.966031, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.1088, + "median": 0.106144, + "min": 0.104704, + "p90": 0.106913, + "sample_count": 941 + }, + "host_enqueue_ms": { + "max": 0.097216, + "median": 0.045152, + "min": 0.038208, + "p90": 0.062688, + "sample_count": 941 + }, + "sample_count": 941, + "synchronized_e2e_ms": { + "max": 0.194592, + "median": 0.153888, + "min": 0.147744, + "p90": 0.168, + "sample_count": 941 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.064609, + "submission_ms": 0.064609, + "synchronized_e2e_ms": 0.253185 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.043264, + "submission_ms": 0.043264, + "synchronized_e2e_ms": 0.233312 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.232288, + "submission_ms": 0.232288, + "synchronized_e2e_ms": 0.419649 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 6.955047, + "submission_ms": 6.955047, + "synchronized_e2e_ms": 7.135847 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.07888, + "submission_ms": 0.07888, + "synchronized_e2e_ms": 0.107552 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.04848, + "submission_ms": 0.04848, + "synchronized_e2e_ms": 0.099712 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.359553, + "submission_ms": 1.359553, + "synchronized_e2e_ms": 1.380033 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.208033, + "submission_ms": 1.208033, + "synchronized_e2e_ms": 1.227425 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.620353, + "submission_ms": 0.620353, + "synchronized_e2e_ms": 0.705057 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 75.892303, + "submission_ms": 75.892303, + "synchronized_e2e_ms": 75.966031 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 2.602461, + "evolution_kernel_ms": 0.302655, + "evolution_speedup": 8.5988, + "evolution_tflops": 39.0252, + "expected_route": "d352_exactd_splitk_c95c_v2", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_1d49_d352_request_highk_b2_n1024_k8192_d352", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 4268, + "measurement_schedule_sha256": "82593d8fc9e109ef043c1ee701c0aa553cae7c930bd6109567a651628219b233", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 1193, + "public_pair_count": 941, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 1193, + "baseline_public_raw": 941, + "candidate_precomputed": 1193, + "candidate_public_raw": 941 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 856 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:adjacent_1d49_d352_request_highk_b2_n1024_k8192_d352", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 2.946784073506891, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.242260605115409, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.09393470879109111, + "including_init_synchronized_e2e_speedup": 0.07915463675569828, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.13240109313098547, + "including_init_synchronized_e2e_speedup": 0.08485631766659645, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.4528077083588927, + "including_init_synchronized_e2e_speedup": 0.14026240136464188, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.5317700085193242, + "including_init_synchronized_e2e_speedup": 0.5697847557584513, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.3433825746156165, + "hot_synchronized_e2e_speedup": 2.242260605115409, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 1493521, + "selected_route": "d352_exactd_splitk_c95c_v2", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_1d49_d352_request_highk_b2_n1024_k8192_d352", + "source": "request_specific", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 416, + "K": 512, + "N": 3840, + "baseline_07cf_adapter_bench_iters": 1972, + "baseline_07cf_adapter_gpu_span_ms": 0.080928, + "baseline_07cf_adapter_host_enqueue_ms": 0.15488049999999998, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.018656, + "baseline_07cf_adapter_kernel_sum_ms": 0.062176, + "baseline_07cf_adapter_submission_ms": 0.15488049999999998, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.184832, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.062176 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.194273, + "synchronized_e2e_ms": 0.222113 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.080928 + }, + "host_enqueue_ms": { + "median": 0.15488049999999998 + }, + "inter_kernel_gap_ms": { + "median": 0.018656 + }, + "kernel_sum_ms": { + "median": 0.062176 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1972, + "submission_ms": { + "median": 0.15488049999999998 + }, + "synchronized_e2e_ms": { + "median": 0.184832 + } + }, + "baseline_07cf_precomputed_bench_iters": 3603, + "baseline_07cf_precomputed_gpu_span_ms": 0.041056, + "baseline_07cf_precomputed_host_enqueue_ms": 0.044576, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.041056, + "baseline_07cf_precomputed_submission_ms": 0.044576, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.091008, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.041056 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.04496, + "synchronized_e2e_ms": 0.082304 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.041056 + }, + "host_enqueue_ms": { + "median": 0.044576 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.041056 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3603, + "submission_ms": { + "median": 0.044576 + }, + "synchronized_e2e_ms": { + "median": 0.091008 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.97116134060795, + "submission": 3.474526651112706, + "synchronized_e2e": 2.030942334739803 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 6.246822, + "after_init_synchronized_e2e_ms_per_call": 6.27271, + "including_init_host_enqueue_ms_per_call": 40.696266, + "including_init_synchronized_e2e_ms_per_call": 490.314204, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.76407465, + "after_init_synchronized_e2e_ms_per_call": 0.7936198, + "including_init_host_enqueue_ms_per_call": 4.20901905, + "including_init_synchronized_e2e_ms_per_call": 49.197769199999996, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.21579991499999998, + "after_init_synchronized_e2e_ms_per_call": 0.24571078, + "including_init_host_enqueue_ms_per_call": 0.560294355, + "including_init_synchronized_e2e_ms_per_call": 5.08612572, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.16097244149999998, + "after_init_synchronized_e2e_ms_per_call": 0.190919878, + "including_init_host_enqueue_ms_per_call": 0.19542188549999998, + "including_init_synchronized_e2e_ms_per_call": 0.674961372, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 6.246822, + "synchronized_e2e_ms": 6.27271, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.357505, + "median": 0.080928, + "min": 0.070944, + "p90": 0.10245440000000004, + "sample_count": 1972 + }, + "host_enqueue_ms": { + "max": 0.755168, + "median": 0.15488049999999998, + "min": 0.12688, + "p90": 0.20829120000000007, + "sample_count": 1972 + }, + "sample_count": 1972, + "synchronized_e2e_ms": { + "max": 0.834624, + "median": 0.184832, + "min": 0.159296, + "p90": 0.234912, + "sample_count": 1972 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3603, + "candidate_precomputed_gpu_span_ms": 0.026304, + "candidate_precomputed_host_enqueue_ms": 0.054368, + "candidate_precomputed_inter_kernel_gap_ms": 0.002016, + "candidate_precomputed_kernel_sum_ms": 0.024288, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.054368, + "candidate_precomputed_synchronized_e2e_ms": 0.0664, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.024288 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.041888, + "synchronized_e2e_ms": 0.057024 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.026304 + }, + "host_enqueue_ms": { + "median": 0.054368 + }, + "inter_kernel_gap_ms": { + "median": 0.002016 + }, + "kernel_sum_ms": { + "median": 0.024288 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3603, + "submission_ms": { + "median": 0.054368 + }, + "synchronized_e2e_ms": { + "median": 0.0664 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282ad520", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc045e90a0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.9294403892944039, + "submission": 0.9008240141259565, + "synchronized_e2e": 1.5089156626506024 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 1972, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.050752, + "candidate_public_raw_host_enqueue_ms": 0.048976, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.050976, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.048976, + "candidate_public_raw_synchronized_e2e_ms": 0.100192, + "candidate_public_raw_tflops_from_gpu_span": 128.9232786885246, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.050752 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.07808, + "synchronized_e2e_ms": 0.11408 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.050752 + }, + "host_enqueue_ms": { + "median": 0.048976 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.050976 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1972, + "submission_ms": { + "median": 0.048976 + }, + "synchronized_e2e_ms": { + "median": 0.100192 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 70.586953, + "after_init_synchronized_e2e_ms_per_call": 70.616937, + "including_init_host_enqueue_ms_per_call": 105.46103699999999, + "including_init_synchronized_e2e_ms_per_call": 105.560013, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 7.102773699999998, + "after_init_synchronized_e2e_ms_per_call": 7.1518665, + "including_init_host_enqueue_ms_per_call": 10.590182099999998, + "including_init_synchronized_e2e_ms_per_call": 10.6461741, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.75435577, + "after_init_synchronized_e2e_ms_per_call": 0.80535945, + "including_init_host_enqueue_ms_per_call": 1.10309661, + "including_init_synchronized_e2e_ms_per_call": 1.15479021, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.11951397699999998, + "after_init_synchronized_e2e_ms_per_call": 0.170708745, + "including_init_host_enqueue_ms_per_call": 0.154388061, + "including_init_synchronized_e2e_ms_per_call": 0.20565182099999998, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 70.586953, + "synchronized_e2e_ms": 70.616937, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.051744, + "median": 0.050752, + "min": 0.050176, + "p90": 0.051072, + "sample_count": 1972 + }, + "host_enqueue_ms": { + "max": 14.913551, + "median": 0.048976, + "min": 0.037984, + "p90": 0.0755200000000001, + "sample_count": 1972 + }, + "sample_count": 1972, + "synchronized_e2e_ms": { + "max": 14.946991, + "median": 0.100192, + "min": 0.09024, + "p90": 0.12201050000000001, + "sample_count": 1972 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.06768, + "submission_ms": 0.06768, + "synchronized_e2e_ms": 0.10272 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.04496, + "submission_ms": 0.04496, + "synchronized_e2e_ms": 0.082304 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.194273, + "submission_ms": 0.194273, + "synchronized_e2e_ms": 0.222113 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 6.246822, + "submission_ms": 6.246822, + "synchronized_e2e_ms": 6.27271 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.072288, + "submission_ms": 0.072288, + "synchronized_e2e_ms": 0.088448 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.041888, + "submission_ms": 0.041888, + "synchronized_e2e_ms": 0.057024 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.274209, + "submission_ms": 1.274209, + "synchronized_e2e_ms": 1.297473 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.067777, + "submission_ms": 1.067777, + "synchronized_e2e_ms": 1.090113 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.07808, + "submission_ms": 0.07808, + "synchronized_e2e_ms": 0.11408 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 70.586953, + "submission_ms": 70.586953, + "synchronized_e2e_ms": 70.616937 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.385919, + "evolution_kernel_ms": 0.184, + "evolution_speedup": 2.0974, + "evolution_tflops": 35.5604, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_1d49_d416_random_b4_n3840_k512_d416", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 11150, + "measurement_schedule_sha256": "159a7fa078cdec3a5de0e37554f2567880aeed7901af1389f04a28716604fc04", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3603, + "public_pair_count": 1972, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3603, + "baseline_public_raw": 1972, + "candidate_precomputed": 3603, + "candidate_public_raw": 1972 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2232 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:adjacent_1d49_d416_random_b4_n3840_k512_d416", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.5608272506082725, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.8447780261897158, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.08882727383092247, + "including_init_synchronized_e2e_speedup": 4.644885786438848, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.11096680845482784, + "including_init_synchronized_e2e_speedup": 4.621168951200977, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.3050945512590682, + "including_init_synchronized_e2e_speedup": 4.404372046070602, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.118395416708148, + "including_init_synchronized_e2e_speedup": 3.28205881532165, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.594577553593947, + "hot_synchronized_e2e_speedup": 1.8447780261897158, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 1494161, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_1d49_d416_random_b4_n3840_k512_d416", + "source": "random_legal", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 480, + "K": 256, + "N": 2816, + "baseline_07cf_adapter_bench_iters": 3590, + "baseline_07cf_adapter_gpu_span_ms": 0.055968, + "baseline_07cf_adapter_host_enqueue_ms": 0.150432, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.033152, + "baseline_07cf_adapter_kernel_sum_ms": 0.022752, + "baseline_07cf_adapter_submission_ms": 0.150432, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.16859200000000002, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.022752 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.340352, + "synchronized_e2e_ms": 0.367872 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.055968 + }, + "host_enqueue_ms": { + "median": 0.150432 + }, + "inter_kernel_gap_ms": { + "median": 0.033152 + }, + "kernel_sum_ms": { + "median": 0.022752 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3590, + "submission_ms": { + "median": 0.150432 + }, + "synchronized_e2e_ms": { + "median": 0.16859200000000002 + } + }, + "baseline_07cf_precomputed_bench_iters": 7336, + "baseline_07cf_precomputed_gpu_span_ms": 0.013568, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042816, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.013568, + "baseline_07cf_precomputed_submission_ms": 0.042816, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.0617125, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.013568 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.069408, + "synchronized_e2e_ms": 0.093408 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.013568 + }, + "host_enqueue_ms": { + "median": 0.042816 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.013568 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 7336, + "submission_ms": { + "median": 0.042816 + }, + "synchronized_e2e_ms": { + "median": 0.0617125 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 4.125, + "submission": 3.5134529147982065, + "synchronized_e2e": 2.7318938626696374 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.6266, + "after_init_synchronized_e2e_ms_per_call": 7.65284, + "including_init_host_enqueue_ms_per_call": 43.177868000000004, + "including_init_synchronized_e2e_ms_per_call": 43.284076999999996, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8980488, + "after_init_synchronized_e2e_ms_per_call": 0.9170168000000001, + "including_init_host_enqueue_ms_per_call": 4.4531756, + "including_init_synchronized_e2e_ms_per_call": 4.480140499999999, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22519368, + "after_init_synchronized_e2e_ms_per_call": 0.24343448, + "including_init_host_enqueue_ms_per_call": 0.5807063600000001, + "including_init_synchronized_e2e_ms_per_call": 0.5997468499999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15790816800000002, + "after_init_synchronized_e2e_ms_per_call": 0.17607624800000002, + "including_init_host_enqueue_ms_per_call": 0.19345943600000004, + "including_init_synchronized_e2e_ms_per_call": 0.21170748500000003, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.6266, + "synchronized_e2e_ms": 7.65284, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 38.903144, + "median": 0.055968, + "min": 0.045408, + "p90": 0.079104, + "sample_count": 3590 + }, + "host_enqueue_ms": { + "max": 39.239624, + "median": 0.150432, + "min": 0.125792, + "p90": 0.2280416, + "sample_count": 3590 + }, + "sample_count": 3590, + "synchronized_e2e_ms": { + "max": 39.328233, + "median": 0.16859200000000002, + "min": 0.140544, + "p90": 0.2515712, + "sample_count": 3590 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 7336, + "candidate_precomputed_gpu_span_ms": 0.019488, + "candidate_precomputed_host_enqueue_ms": 0.053248, + "candidate_precomputed_inter_kernel_gap_ms": 0.003072, + "candidate_precomputed_kernel_sum_ms": 0.01648, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.053248, + "candidate_precomputed_synchronized_e2e_ms": 0.0648, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.01648 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.062368, + "synchronized_e2e_ms": 0.08368 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.019488 + }, + "host_enqueue_ms": { + "median": 0.053248 + }, + "inter_kernel_gap_ms": { + "median": 0.003072 + }, + "kernel_sum_ms": { + "median": 0.01648 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 7336, + "submission_ms": { + "median": 0.053248 + }, + "synchronized_e2e_ms": { + "median": 0.0648 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295c2d20", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295c2930" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.4449917898193763, + "submission": 0.8497596153846154, + "synchronized_e2e": 1.1466666666666667 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 3590, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.02816, + "candidate_public_raw_host_enqueue_ms": 0.045248, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.028416, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.045248, + "candidate_public_raw_synchronized_e2e_ms": 0.074304, + "candidate_public_raw_tflops_from_gpu_span": 49.152, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.02816 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.0888, + "synchronized_e2e_ms": 0.115168 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.02816 + }, + "host_enqueue_ms": { + "median": 0.045248 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.028416 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3590, + "submission_ms": { + "median": 0.045248 + }, + "synchronized_e2e_ms": { + "median": 0.074304 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 44.795055, + "after_init_synchronized_e2e_ms_per_call": 44.834255, + "including_init_host_enqueue_ms_per_call": 80.63307499999999, + "including_init_synchronized_e2e_ms_per_call": 528.934468, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 4.5202287, + "after_init_synchronized_e2e_ms_per_call": 4.5502991, + "including_init_host_enqueue_ms_per_call": 8.1040307, + "including_init_synchronized_e2e_ms_per_call": 52.9603204, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.49274607, + "after_init_synchronized_e2e_ms_per_call": 0.52190351, + "including_init_host_enqueue_ms_per_call": 0.8511262699999999, + "including_init_synchronized_e2e_ms_per_call": 5.36290564, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.089997807, + "after_init_synchronized_e2e_ms_per_call": 0.11906395099999999, + "including_init_host_enqueue_ms_per_call": 0.12583582699999998, + "including_init_synchronized_e2e_ms_per_call": 0.6031641640000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 44.795055, + "synchronized_e2e_ms": 44.834255, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.028896, + "median": 0.02816, + "min": 0.02752, + "p90": 0.028448, + "sample_count": 3590 + }, + "host_enqueue_ms": { + "max": 49.446964, + "median": 0.045248, + "min": 0.036, + "p90": 0.07158719999999999, + "sample_count": 3590 + }, + "sample_count": 3590, + "synchronized_e2e_ms": { + "max": 49.578004, + "median": 0.074304, + "min": 0.066624, + "p90": 0.09642239999999999, + "sample_count": 3590 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.10048, + "submission_ms": 0.10048, + "synchronized_e2e_ms": 0.12896 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.069408, + "submission_ms": 0.069408, + "synchronized_e2e_ms": 0.093408 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.340352, + "submission_ms": 0.340352, + "synchronized_e2e_ms": 0.367872 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.6266, + "submission_ms": 7.6266, + "synchronized_e2e_ms": 7.65284 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.10592, + "submission_ms": 0.10592, + "synchronized_e2e_ms": 0.132128 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.062368, + "submission_ms": 0.062368, + "synchronized_e2e_ms": 0.08368 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.684002, + "submission_ms": 1.684002, + "synchronized_e2e_ms": 1.721058 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.55053, + "submission_ms": 1.55053, + "synchronized_e2e_ms": 1.585058 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.0888, + "submission_ms": 0.0888, + "synchronized_e2e_ms": 0.115168 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 44.795055, + "submission_ms": 44.795055, + "synchronized_e2e_ms": 44.834255 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.332511, + "evolution_kernel_ms": 0.178015, + "evolution_speedup": 1.8679, + "evolution_tflops": 7.7753, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_1d49_d480_smallk_boundary_b2_n2816_k256_d480", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 21852, + "measurement_schedule_sha256": "284cdf4678a209fc2bce61709fe3684b319fe4d2a48e30e41c90ca746f9ed57b", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 7336, + "public_pair_count": 3590, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 7336, + "baseline_public_raw": 3590, + "candidate_precomputed": 7336, + "candidate_public_raw": 3590 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 4372 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:adjacent_1d49_d480_smallk_boundary_b2_n2816_k256_d480", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.6962233169129721, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.2689491817398797, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.17069180696768577, + "including_init_synchronized_e2e_speedup": 0.08183258913654308, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.20152890608883273, + "including_init_synchronized_e2e_speedup": 0.08459428617807227, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.4664357976822191, + "including_init_synchronized_e2e_speedup": 0.111832444995247, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.4788376038352704, + "including_init_synchronized_e2e_speedup": 0.3509947998170528, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.9874999999999998, + "hot_synchronized_e2e_speedup": 2.2689491817398797, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 1494801, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_1d49_d480_smallk_boundary_b2_n2816_k256_d480", + "source": "guard_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 48, + "K": 256, + "N": 3968, + "baseline_07cf_adapter_bench_iters": 5153, + "baseline_07cf_adapter_gpu_span_ms": 0.056064, + "baseline_07cf_adapter_host_enqueue_ms": 0.160736, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.038208, + "baseline_07cf_adapter_kernel_sum_ms": 0.017824, + "baseline_07cf_adapter_submission_ms": 0.160736, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.18016, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.017824 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.340225, + "synchronized_e2e_ms": 0.362785 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.056064 + }, + "host_enqueue_ms": { + "median": 0.160736 + }, + "inter_kernel_gap_ms": { + "median": 0.038208 + }, + "kernel_sum_ms": { + "median": 0.017824 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 5153, + "submission_ms": { + "median": 0.160736 + }, + "synchronized_e2e_ms": { + "median": 0.18016 + } + }, + "baseline_07cf_precomputed_bench_iters": 13186, + "baseline_07cf_precomputed_gpu_span_ms": 0.007552, + "baseline_07cf_precomputed_host_enqueue_ms": 0.045232499999999995, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.007552, + "baseline_07cf_precomputed_submission_ms": 0.045232499999999995, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.060128, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.007552 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.054368, + "synchronized_e2e_ms": 0.070656 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.007552 + }, + "host_enqueue_ms": { + "median": 0.045232499999999995 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.007552 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 13186, + "submission_ms": { + "median": 0.045232499999999995 + }, + "synchronized_e2e_ms": { + "median": 0.060128 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 7.423728813559323, + "submission": 3.5535510971093793, + "synchronized_e2e": 2.996274614156466 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.025769, + "after_init_synchronized_e2e_ms_per_call": 8.053417, + "including_init_host_enqueue_ms_per_call": 44.960207, + "including_init_synchronized_e2e_ms_per_call": 463.16473600000006, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9472393, + "after_init_synchronized_e2e_ms_per_call": 0.9674856999999999, + "including_init_host_enqueue_ms_per_call": 4.6406830999999995, + "including_init_synchronized_e2e_ms_per_call": 46.47861760000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23938632999999998, + "after_init_synchronized_e2e_ms_per_call": 0.25889257, + "including_init_host_enqueue_ms_per_call": 0.60873071, + "including_init_synchronized_e2e_ms_per_call": 4.810005760000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.168601033, + "after_init_synchronized_e2e_ms_per_call": 0.18803325699999998, + "including_init_host_enqueue_ms_per_call": 0.205535471, + "including_init_synchronized_e2e_ms_per_call": 0.6431445760000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.025769, + "synchronized_e2e_ms": 8.053417, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 52.302171, + "median": 0.056064, + "min": 0.044192, + "p90": 0.07984, + "sample_count": 5153 + }, + "host_enqueue_ms": { + "max": 52.823543, + "median": 0.160736, + "min": 0.128896, + "p90": 0.21257040000000002, + "sample_count": 5153 + }, + "sample_count": 5153, + "synchronized_e2e_ms": { + "max": 52.907191, + "median": 0.18016, + "min": 0.145888, + "p90": 0.23606400000000002, + "sample_count": 5153 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 13186, + "candidate_precomputed_gpu_span_ms": 0.014624, + "candidate_precomputed_host_enqueue_ms": 0.05728, + "candidate_precomputed_inter_kernel_gap_ms": 0.005984, + "candidate_precomputed_kernel_sum_ms": 0.008736, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.05728, + "candidate_precomputed_synchronized_e2e_ms": 0.069472, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.008736 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.058656, + "synchronized_e2e_ms": 0.074016 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.014624 + }, + "host_enqueue_ms": { + "median": 0.05728 + }, + "inter_kernel_gap_ms": { + "median": 0.005984 + }, + "kernel_sum_ms": { + "median": 0.008736 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 13186, + "submission_ms": { + "median": 0.05728 + }, + "synchronized_e2e_ms": { + "median": 0.069472 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295c1d90", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295c30e0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.3457330415754922, + "submission": 0.8731843575418995, + "synchronized_e2e": 1.0285582680792262 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 5153, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.01968, + "candidate_public_raw_host_enqueue_ms": 0.050016, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.019872, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.050016, + "candidate_public_raw_synchronized_e2e_ms": 0.071456, + "candidate_public_raw_tflops_from_gpu_span": 19.820643902439027, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.01968 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.076064, + "synchronized_e2e_ms": 0.096608 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.01968 + }, + "host_enqueue_ms": { + "median": 0.050016 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.019872 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 5153, + "submission_ms": { + "median": 0.050016 + }, + "synchronized_e2e_ms": { + "median": 0.071456 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 83.722038, + "after_init_synchronized_e2e_ms_per_call": 83.779798, + "including_init_host_enqueue_ms_per_call": 121.080061, + "including_init_synchronized_e2e_ms_per_call": 121.21542099999999, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 8.417218199999999, + "after_init_synchronized_e2e_ms_per_call": 8.442290199999999, + "including_init_host_enqueue_ms_per_call": 12.1530205, + "including_init_synchronized_e2e_ms_per_call": 12.1858525, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.88673622, + "after_init_synchronized_e2e_ms_per_call": 0.90853942, + "including_init_host_enqueue_ms_per_call": 1.26031645, + "including_init_synchronized_e2e_ms_per_call": 1.2828956499999997, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.133688022, + "after_init_synchronized_e2e_ms_per_call": 0.155164342, + "including_init_host_enqueue_ms_per_call": 0.17104604499999998, + "including_init_synchronized_e2e_ms_per_call": 0.19259996499999998, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 83.722038, + "synchronized_e2e_ms": 83.779798, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.020288, + "median": 0.01968, + "min": 0.019232, + "p90": 0.019936, + "sample_count": 5153 + }, + "host_enqueue_ms": { + "max": 15.636368, + "median": 0.050016, + "min": 0.039168, + "p90": 0.07363200000000003, + "sample_count": 5153 + }, + "sample_count": 5153, + "synchronized_e2e_ms": { + "max": 78.728049, + "median": 0.071456, + "min": 0.062176, + "p90": 0.09439360000000001, + "sample_count": 5153 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.084672, + "submission_ms": 0.084672, + "synchronized_e2e_ms": 0.102496 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.054368, + "submission_ms": 0.054368, + "synchronized_e2e_ms": 0.070656 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.340225, + "submission_ms": 0.340225, + "synchronized_e2e_ms": 0.362785 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.025769, + "submission_ms": 8.025769, + "synchronized_e2e_ms": 8.053417 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.098464, + "submission_ms": 0.098464, + "synchronized_e2e_ms": 0.116768 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.058656, + "submission_ms": 0.058656, + "synchronized_e2e_ms": 0.074016 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.592386, + "submission_ms": 1.592386, + "synchronized_e2e_ms": 1.618338 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.177921, + "submission_ms": 1.177921, + "synchronized_e2e_ms": 1.202625 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.076064, + "submission_ms": 0.076064, + "synchronized_e2e_ms": 0.096608 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 83.722038, + "submission_ms": 83.722038, + "synchronized_e2e_ms": 83.779798 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.228768, + "evolution_kernel_ms": 0.17192, + "evolution_speedup": 1.3307, + "evolution_tflops": 2.2689, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_1d49_d48_lowk_boundary_b4_n3968_k256_d48", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 36678, + "measurement_schedule_sha256": "45a1fb56f62253889f8a6c18fb3c10cd2586d9e2ba27a791714953ff18df9179", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 13186, + "public_pair_count": 5153, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 13186, + "baseline_public_raw": 5153, + "candidate_precomputed": 13186, + "candidate_public_raw": 5153 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 7338 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:adjacent_1d49_d48_lowk_boundary_b4_n3968_k256_d48", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.5164113785557987, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.5212718316166587, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.09612600164063417, + "including_init_synchronized_e2e_speedup": 3.8210050518242236, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.11459991034186436, + "including_init_synchronized_e2e_speedup": 3.81414575631865, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.28495469134404755, + "including_init_synchronized_e2e_speedup": 3.7493351544219533, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.2118329158383565, + "including_init_synchronized_e2e_speedup": 3.339276702360772, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.8487804878048784, + "hot_synchronized_e2e_speedup": 2.5212718316166587, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 1490481, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_1d49_d48_lowk_boundary_b4_n3968_k256_d48", + "source": "guard_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 64, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 112, + "K": 8192, + "N": 512, + "baseline_07cf_adapter_bench_iters": 1225, + "baseline_07cf_adapter_gpu_span_ms": 0.131136, + "baseline_07cf_adapter_host_enqueue_ms": 0.140672, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.03024, + "baseline_07cf_adapter_kernel_sum_ms": 0.100928, + "baseline_07cf_adapter_submission_ms": 0.140672, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.227777, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.100928 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.21504, + "synchronized_e2e_ms": 0.293568 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.131136 + }, + "host_enqueue_ms": { + "median": 0.140672 + }, + "inter_kernel_gap_ms": { + "median": 0.03024 + }, + "kernel_sum_ms": { + "median": 0.100928 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1225, + "submission_ms": { + "median": 0.140672 + }, + "synchronized_e2e_ms": { + "median": 0.227777 + } + }, + "baseline_07cf_precomputed_bench_iters": 1423, + "baseline_07cf_precomputed_gpu_span_ms": 0.106176, + "baseline_07cf_precomputed_host_enqueue_ms": 0.041632, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.106176, + "baseline_07cf_precomputed_submission_ms": 0.041632, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.154529, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.106176 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.044928, + "synchronized_e2e_ms": 0.128608 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.106176 + }, + "host_enqueue_ms": { + "median": 0.041632 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.106176 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1423, + "submission_ms": { + "median": 0.041632 + }, + "synchronized_e2e_ms": { + "median": 0.154529 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.2350813743218807, + "submission": 3.378939277478862, + "synchronized_e2e": 1.474008114981654 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 6.966407, + "after_init_synchronized_e2e_ms_per_call": 7.034247, + "including_init_host_enqueue_ms_per_call": 41.194314000000006, + "including_init_synchronized_e2e_ms_per_call": 41.379690000000004, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8232455, + "after_init_synchronized_e2e_ms_per_call": 0.9084239999999999, + "including_init_host_enqueue_ms_per_call": 4.246036200000001, + "including_init_synchronized_e2e_ms_per_call": 4.342968300000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.20892935, + "after_init_synchronized_e2e_ms_per_call": 0.2958417, + "including_init_host_enqueue_ms_per_call": 0.55120842, + "including_init_synchronized_e2e_ms_per_call": 0.6392961300000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.14749773500000002, + "after_init_synchronized_e2e_ms_per_call": 0.23458347000000002, + "including_init_host_enqueue_ms_per_call": 0.181725642, + "including_init_synchronized_e2e_ms_per_call": 0.26892891300000005, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 6.966407, + "synchronized_e2e_ms": 7.034247, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.309344, + "median": 0.131136, + "min": 0.123745, + "p90": 0.144832, + "sample_count": 1225 + }, + "host_enqueue_ms": { + "max": 0.597792, + "median": 0.140672, + "min": 0.121472, + "p90": 0.1768256000000001, + "sample_count": 1225 + }, + "sample_count": 1225, + "synchronized_e2e_ms": { + "max": 0.626528, + "median": 0.227777, + "min": 0.211905, + "p90": 0.2591104, + "sample_count": 1225 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 1423, + "candidate_precomputed_gpu_span_ms": 0.070496, + "candidate_precomputed_host_enqueue_ms": 0.036896, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.070496, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.036896, + "candidate_precomputed_synchronized_e2e_ms": 0.104736, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.070496 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.032608, + "synchronized_e2e_ms": 0.099968 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.070496 + }, + "host_enqueue_ms": { + "median": 0.036896 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.070496 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1423, + "submission_ms": { + "median": 0.036896 + }, + "synchronized_e2e_ms": { + "median": 0.104736 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffe4324f170", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04737b90" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.158874262369496, + "submission": 1.0858629661751953, + "synchronized_e2e": 1.1811793461655973 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 1225, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.081696, + "candidate_public_raw_host_enqueue_ms": 0.040064, + "candidate_public_raw_inter_kernel_gap_ms": 0.000192, + "candidate_public_raw_kernel_sum_ms": 0.081504, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.040064, + "candidate_public_raw_synchronized_e2e_ms": 0.123712, + "candidate_public_raw_tflops_from_gpu_span": 23.0004919702311, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.081504 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.067681, + "synchronized_e2e_ms": 0.136353 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.081696 + }, + "host_enqueue_ms": { + "median": 0.040064 + }, + "inter_kernel_gap_ms": { + "median": 0.000192 + }, + "kernel_sum_ms": { + "median": 0.081504 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1225, + "submission_ms": { + "median": 0.040064 + }, + "synchronized_e2e_ms": { + "median": 0.123712 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 43.858157, + "after_init_synchronized_e2e_ms_per_call": 43.919213, + "including_init_host_enqueue_ms_per_call": 78.435697, + "including_init_synchronized_e2e_ms_per_call": 492.007005, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 4.4218733, + "after_init_synchronized_e2e_ms_per_call": 4.5032621, + "including_init_host_enqueue_ms_per_call": 7.8796273, + "including_init_synchronized_e2e_ms_per_call": 49.3120413, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.47824492999999996, + "after_init_synchronized_e2e_ms_per_call": 0.5616670100000001, + "including_init_host_enqueue_ms_per_call": 0.8240203300000001, + "including_init_synchronized_e2e_ms_per_call": 5.04254493, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.08388209299999999, + "after_init_synchronized_e2e_ms_per_call": 0.167507501, + "including_init_host_enqueue_ms_per_call": 0.118459633, + "including_init_synchronized_e2e_ms_per_call": 0.6155952929999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 43.858157, + "synchronized_e2e_ms": 43.919213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.084672, + "median": 0.081696, + "min": 0.080032, + "p90": 0.082816, + "sample_count": 1225 + }, + "host_enqueue_ms": { + "max": 0.118944, + "median": 0.040064, + "min": 0.031872, + "p90": 0.05536640000000001, + "sample_count": 1225 + }, + "sample_count": 1225, + "synchronized_e2e_ms": { + "max": 0.190144, + "median": 0.123712, + "min": 0.11648, + "p90": 0.1374276, + "sample_count": 1225 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.064704, + "submission_ms": 0.064704, + "synchronized_e2e_ms": 0.14784 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.044928, + "submission_ms": 0.044928, + "synchronized_e2e_ms": 0.128608 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.21504, + "submission_ms": 0.21504, + "synchronized_e2e_ms": 0.293568 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 6.966407, + "submission_ms": 6.966407, + "synchronized_e2e_ms": 7.034247 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.054912, + "submission_ms": 0.054912, + "synchronized_e2e_ms": 0.115904 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.032608, + "submission_ms": 0.032608, + "synchronized_e2e_ms": 0.099968 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 0.978529, + "submission_ms": 0.978529, + "synchronized_e2e_ms": 1.001153 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 0.878081, + "submission_ms": 0.878081, + "synchronized_e2e_ms": 0.898657 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.067681, + "submission_ms": 0.067681, + "synchronized_e2e_ms": 0.136353 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 43.858157, + "submission_ms": 43.858157, + "synchronized_e2e_ms": 43.919213 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 1.758429, + "evolution_kernel_ms": 0.295872, + "evolution_speedup": 5.9432, + "evolution_tflops": 6.3509, + "expected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_3328_d112_highk_low_n_b2_n512_k8192_d112", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 5296, + "measurement_schedule_sha256": "a82aedca3715134dc8aa6422a754f79281b495e3f2e62b7440d5ee8347cb71f2", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 1423, + "public_pair_count": 1225, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 1423, + "baseline_public_raw": 1225, + "candidate_precomputed": 1423, + "candidate_public_raw": 1225 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 1060 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:adjacent_3328_d112_highk_low_n_b2_n512_k8192_d112", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.5061280072628234, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.8411875969994826, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.1601633207771733, + "including_init_synchronized_e2e_speedup": 0.08410386352121146, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.2017257667502853, + "including_init_synchronized_e2e_speedup": 0.08807115230900005, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.5267208056246706, + "including_init_synchronized_e2e_speedup": 0.12678045290119014, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.4004356139251342, + "including_init_synchronized_e2e_speedup": 0.4368599241385039, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.6051703877790833, + "hot_synchronized_e2e_speedup": 1.8411875969994826, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 3328112, + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_3328_d112_highk_low_n_b2_n512_k8192_d112", + "source": "heldout_neighborhood", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 128, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 4, + "num_warps": 4 + } + }, + { + "B": 8, + "D": 128, + "K": 512, + "N": 8192, + "baseline_07cf_adapter_bench_iters": 1668, + "baseline_07cf_adapter_gpu_span_ms": 0.0750405, + "baseline_07cf_adapter_host_enqueue_ms": 0.154016, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.006592, + "baseline_07cf_adapter_kernel_sum_ms": 0.068416, + "baseline_07cf_adapter_submission_ms": 0.154016, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.178144, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.068416 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.201248, + "synchronized_e2e_ms": 0.222912 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.0750405 + }, + "host_enqueue_ms": { + "median": 0.154016 + }, + "inter_kernel_gap_ms": { + "median": 0.006592 + }, + "kernel_sum_ms": { + "median": 0.068416 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1668, + "submission_ms": { + "median": 0.154016 + }, + "synchronized_e2e_ms": { + "median": 0.178144 + } + }, + "baseline_07cf_precomputed_bench_iters": 4801, + "baseline_07cf_precomputed_gpu_span_ms": 0.029088, + "baseline_07cf_precomputed_host_enqueue_ms": 0.043168, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.029088, + "baseline_07cf_precomputed_submission_ms": 0.043168, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.077984, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.029088 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.047488, + "synchronized_e2e_ms": 0.07888 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.029088 + }, + "host_enqueue_ms": { + "median": 0.043168 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.029088 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4801, + "submission_ms": { + "median": 0.043168 + }, + "synchronized_e2e_ms": { + "median": 0.077984 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.5797751650165015, + "submission": 3.5678280207561155, + "synchronized_e2e": 2.2843660237997536 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 6.731463, + "after_init_synchronized_e2e_ms_per_call": 6.757543, + "including_init_host_enqueue_ms_per_call": 41.180907, + "including_init_synchronized_e2e_ms_per_call": 490.799037, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8117607, + "after_init_synchronized_e2e_ms_per_call": 0.8360839, + "including_init_host_enqueue_ms_per_call": 4.2567051, + "including_init_synchronized_e2e_ms_per_call": 49.2402333, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.21979047, + "after_init_synchronized_e2e_ms_per_call": 0.24393799000000002, + "including_init_host_enqueue_ms_per_call": 0.5642849099999999, + "including_init_synchronized_e2e_ms_per_call": 5.08435293, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.16059344699999997, + "after_init_synchronized_e2e_ms_per_call": 0.184723399, + "including_init_host_enqueue_ms_per_call": 0.19504289099999997, + "including_init_synchronized_e2e_ms_per_call": 0.6687648930000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 6.731463, + "synchronized_e2e_ms": 6.757543, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.317184, + "median": 0.0750405, + "min": 0.068672, + "p90": 0.09664, + "sample_count": 1668 + }, + "host_enqueue_ms": { + "max": 40.374794, + "median": 0.154016, + "min": 0.131488, + "p90": 0.2039111, + "sample_count": 1668 + }, + "sample_count": 1668, + "synchronized_e2e_ms": { + "max": 40.45649, + "median": 0.178144, + "min": 0.15888, + "p90": 0.228128, + "sample_count": 1668 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 4801, + "candidate_precomputed_gpu_span_ms": 0.020864, + "candidate_precomputed_host_enqueue_ms": 0.040128, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.020864, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.040128, + "candidate_precomputed_synchronized_e2e_ms": 0.0544, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.020864 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.034112, + "synchronized_e2e_ms": 0.049696 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.020864 + }, + "host_enqueue_ms": { + "median": 0.040128 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.020864 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4801, + "submission_ms": { + "median": 0.040128 + }, + "synchronized_e2e_ms": { + "median": 0.0544 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7f305af0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7f304d40" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.880368098159509, + "submission": 1.20933014354067, + "synchronized_e2e": 1.9900000000000002 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 1668, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.060096, + "candidate_public_raw_host_enqueue_ms": 0.048528, + "candidate_public_raw_inter_kernel_gap_ms": 0.00016, + "candidate_public_raw_kernel_sum_ms": 0.059904, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.048528, + "candidate_public_raw_synchronized_e2e_ms": 0.108256, + "candidate_public_raw_tflops_from_gpu_span": 142.93687752928648, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.059904 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.265953, + "synchronized_e2e_ms": 0.312833 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.060096 + }, + "host_enqueue_ms": { + "median": 0.048528 + }, + "inter_kernel_gap_ms": { + "median": 0.00016 + }, + "kernel_sum_ms": { + "median": 0.059904 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1668, + "submission_ms": { + "median": 0.048528 + }, + "synchronized_e2e_ms": { + "median": 0.108256 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 48.198226, + "after_init_synchronized_e2e_ms_per_call": 48.237554, + "including_init_host_enqueue_ms_per_call": 83.07230999999999, + "including_init_synchronized_e2e_ms_per_call": 83.18063000000001, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 4.863497799999999, + "after_init_synchronized_e2e_ms_per_call": 4.921185800000001, + "including_init_host_enqueue_ms_per_call": 8.350906199999999, + "including_init_synchronized_e2e_ms_per_call": 8.4154934, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.53002498, + "after_init_synchronized_e2e_ms_per_call": 0.58954898, + "including_init_host_enqueue_ms_per_call": 0.8787658199999998, + "including_init_synchronized_e2e_ms_per_call": 0.9389797400000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.09667769799999999, + "after_init_synchronized_e2e_ms_per_call": 0.156385298, + "including_init_host_enqueue_ms_per_call": 0.131551782, + "including_init_synchronized_e2e_ms_per_call": 0.191328374, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 48.198226, + "synchronized_e2e_ms": 48.237554, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.061056, + "median": 0.060096, + "min": 0.059392, + "p90": 0.06047999999999999, + "sample_count": 1668 + }, + "host_enqueue_ms": { + "max": 45.822927, + "median": 0.048528, + "min": 0.039808, + "p90": 0.0723075, + "sample_count": 1668 + }, + "sample_count": 1668, + "synchronized_e2e_ms": { + "max": 45.924719, + "median": 0.108256, + "min": 0.100352, + "p90": 0.1290435, + "sample_count": 1668 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.067776, + "submission_ms": 0.067776, + "synchronized_e2e_ms": 0.099008 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.047488, + "submission_ms": 0.047488, + "synchronized_e2e_ms": 0.07888 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.201248, + "submission_ms": 0.201248, + "synchronized_e2e_ms": 0.222912 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 6.731463, + "submission_ms": 6.731463, + "synchronized_e2e_ms": 6.757543 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.054977, + "submission_ms": 0.054977, + "synchronized_e2e_ms": 0.070689 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.034112, + "submission_ms": 0.034112, + "synchronized_e2e_ms": 0.049696 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.229921, + "submission_ms": 1.229921, + "synchronized_e2e_ms": 1.274273 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.098977, + "submission_ms": 1.098977, + "synchronized_e2e_ms": 1.119873 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.265953, + "submission_ms": 0.265953, + "synchronized_e2e_ms": 0.312833 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 48.198226, + "submission_ms": 48.198226, + "synchronized_e2e_ms": 48.237554 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.394592, + "evolution_kernel_ms": 0.177472, + "evolution_speedup": 2.2234, + "evolution_tflops": 48.4016, + "expected_route": "paired_large_v15", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_3328_d128_exact_guard_k_neighbor_b8_n8192_k512_d128", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 12938, + "measurement_schedule_sha256": "f8ec643779dbe46019e579df7f48c7a76bdf38ff698993a565036653bb9e199e", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 4801, + "public_pair_count": 1668, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 4801, + "baseline_public_raw": 1668, + "candidate_precomputed": 4801, + "candidate_public_raw": 1668 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2590 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:adjacent_3328_d128_exact_guard_k_neighbor_b8_n8192_k512_d128", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.3941717791411041, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.645580845403488, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.1400888403255273, + "including_init_synchronized_e2e_speedup": 5.900400574027871, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.16989480462208922, + "including_init_synchronized_e2e_speedup": 5.851140385898229, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.41377052335838155, + "including_init_synchronized_e2e_speedup": 5.414763187542256, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.1812069380076893, + "including_init_synchronized_e2e_speedup": 3.4953774969101032, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.2486771166134185, + "hot_synchronized_e2e_speedup": 1.645580845403488, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 3328128, + "selected_route": "paired_large_v15", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_3328_d128_exact_guard_k_neighbor_b8_n8192_k512_d128", + "source": "guard_miss_fallback", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 3, + "D": 224, + "K": 512, + "N": 3840, + "baseline_07cf_adapter_bench_iters": 3372, + "baseline_07cf_adapter_gpu_span_ms": 0.061184, + "baseline_07cf_adapter_host_enqueue_ms": 0.148144, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.026912, + "baseline_07cf_adapter_kernel_sum_ms": 0.034272, + "baseline_07cf_adapter_submission_ms": 0.148144, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.166192, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.034272 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.173376, + "synchronized_e2e_ms": 0.189856 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.061184 + }, + "host_enqueue_ms": { + "median": 0.148144 + }, + "inter_kernel_gap_ms": { + "median": 0.026912 + }, + "kernel_sum_ms": { + "median": 0.034272 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3372, + "submission_ms": { + "median": 0.148144 + }, + "synchronized_e2e_ms": { + "median": 0.166192 + } + }, + "baseline_07cf_precomputed_bench_iters": 5421, + "baseline_07cf_precomputed_gpu_span_ms": 0.018272, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042304, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.018272, + "baseline_07cf_precomputed_submission_ms": 0.042304, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.066208, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.018272 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.039936, + "synchronized_e2e_ms": 0.06144 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.018272 + }, + "host_enqueue_ms": { + "median": 0.042304 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.018272 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5421, + "submission_ms": { + "median": 0.042304 + }, + "synchronized_e2e_ms": { + "median": 0.066208 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 3.3485113835376534, + "submission": 3.5018910741301057, + "synchronized_e2e": 2.5101498308361525 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 6.909895, + "after_init_synchronized_e2e_ms_per_call": 6.933159, + "including_init_host_enqueue_ms_per_call": 42.461163, + "including_init_synchronized_e2e_ms_per_call": 42.564396, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8243191, + "after_init_synchronized_e2e_ms_per_call": 0.8428886999999999, + "including_init_host_enqueue_ms_per_call": 4.3794458999999994, + "including_init_synchronized_e2e_ms_per_call": 4.4060124, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.21576151, + "after_init_synchronized_e2e_ms_per_call": 0.23386167, + "including_init_host_enqueue_ms_per_call": 0.57127419, + "including_init_synchronized_e2e_ms_per_call": 0.59017404, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.154905751, + "after_init_synchronized_e2e_ms_per_call": 0.172958967, + "including_init_host_enqueue_ms_per_call": 0.190457019, + "including_init_synchronized_e2e_ms_per_call": 0.20859020400000003, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 6.909895, + "synchronized_e2e_ms": 6.933159, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 44.987727, + "median": 0.061184, + "min": 0.052288, + "p90": 0.08584000000000003, + "sample_count": 3372 + }, + "host_enqueue_ms": { + "max": 45.294575, + "median": 0.148144, + "min": 0.121728, + "p90": 0.230112, + "sample_count": 3372 + }, + "sample_count": 3372, + "synchronized_e2e_ms": { + "max": 45.388079, + "median": 0.166192, + "min": 0.138016, + "p90": 0.25441280000000005, + "sample_count": 3372 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 5421, + "candidate_precomputed_gpu_span_ms": 0.015809, + "candidate_precomputed_host_enqueue_ms": 0.052544, + "candidate_precomputed_inter_kernel_gap_ms": 0.002848, + "candidate_precomputed_kernel_sum_ms": 0.01296, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.052544, + "candidate_precomputed_synchronized_e2e_ms": 0.063968, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.01296 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.037376, + "synchronized_e2e_ms": 0.049408 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.015809 + }, + "host_enqueue_ms": { + "median": 0.052544 + }, + "inter_kernel_gap_ms": { + "median": 0.002848 + }, + "kernel_sum_ms": { + "median": 0.01296 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5421, + "submission_ms": { + "median": 0.052544 + }, + "synchronized_e2e_ms": { + "median": 0.063968 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295c3320", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295c32f0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.8824720096147765, + "submission": 0.848355663824604, + "synchronized_e2e": 1.1715857928964482 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 3372, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.02976, + "candidate_public_raw_host_enqueue_ms": 0.044576, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.029728, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.044576, + "candidate_public_raw_synchronized_e2e_ms": 0.074944, + "candidate_public_raw_tflops_from_gpu_span": 88.79070967741936, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.029728 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.045984, + "synchronized_e2e_ms": 0.066848 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.02976 + }, + "host_enqueue_ms": { + "median": 0.044576 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.029728 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3372, + "submission_ms": { + "median": 0.044576 + }, + "synchronized_e2e_ms": { + "median": 0.074944 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.13469, + "after_init_synchronized_e2e_ms_per_call": 2.15805, + "including_init_host_enqueue_ms_per_call": 37.97271, + "including_init_synchronized_e2e_ms_per_call": 486.258263, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.25358739999999996, + "after_init_synchronized_e2e_ms_per_call": 0.28325459999999997, + "including_init_host_enqueue_ms_per_call": 3.8373894, + "including_init_synchronized_e2e_ms_per_call": 48.693275899999996, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06547714, + "after_init_synchronized_e2e_ms_per_call": 0.09577506, + "including_init_host_enqueue_ms_per_call": 0.42385733999999997, + "including_init_synchronized_e2e_ms_per_call": 4.93677719, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.046666113999999995, + "after_init_synchronized_e2e_ms_per_call": 0.077027106, + "including_init_host_enqueue_ms_per_call": 0.08250413399999999, + "including_init_synchronized_e2e_ms_per_call": 0.5611273189999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.13469, + "synchronized_e2e_ms": 2.15805, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.030496, + "median": 0.02976, + "min": 0.029312, + "p90": 0.030016, + "sample_count": 3372 + }, + "host_enqueue_ms": { + "max": 0.417537, + "median": 0.044576, + "min": 0.03504, + "p90": 0.070432, + "sample_count": 3372 + }, + "sample_count": 3372, + "synchronized_e2e_ms": { + "max": 41.202411, + "median": 0.074944, + "min": 0.06688, + "p90": 0.09700800000000001, + "sample_count": 3372 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.062592, + "submission_ms": 0.062592, + "synchronized_e2e_ms": 0.082624 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.039936, + "submission_ms": 0.039936, + "synchronized_e2e_ms": 0.06144 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.173376, + "submission_ms": 0.173376, + "synchronized_e2e_ms": 0.189856 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 6.909895, + "submission_ms": 6.909895, + "synchronized_e2e_ms": 6.933159 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.06688, + "submission_ms": 0.06688, + "synchronized_e2e_ms": 0.081184 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.037376, + "submission_ms": 0.037376, + "synchronized_e2e_ms": 0.049408 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.173409, + "submission_ms": 1.173409, + "synchronized_e2e_ms": 1.199233 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.568802, + "submission_ms": 1.568802, + "synchronized_e2e_ms": 1.592066 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.045984, + "submission_ms": 0.045984, + "synchronized_e2e_ms": 0.066848 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.13469, + "submission_ms": 2.13469, + "synchronized_e2e_ms": 2.15805 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.318207, + "evolution_kernel_ms": 0.174895, + "evolution_speedup": 1.8194, + "evolution_tflops": 15.1086, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_3328_d224_tail_div_b3_n3840_k512_d224", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 17586, + "measurement_schedule_sha256": "7c7c2b49bf2906de90176383821f66d442d452d8b38117cb0ff90c89257679e7", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 5421, + "public_pair_count": 3372, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 5421, + "baseline_public_raw": 3372, + "candidate_precomputed": 5421, + "candidate_public_raw": 3372 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 3520 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:adjacent_3328_d224_tail_div_b3_n3840_k512_d224", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.1557973306344487, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.217549103330487, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.21269618405505, + "including_init_synchronized_e2e_speedup": 0.08753454540267627, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.975728196470596, + "including_init_synchronized_e2e_speedup": 0.09048502731770405, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.4417804593387884, + "including_init_synchronized_e2e_speedup": 0.11954642012109928, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.2454299010013434, + "including_init_synchronized_e2e_speedup": 0.37173418034918393, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.0559139784946234, + "hot_synchronized_e2e_speedup": 2.217549103330487, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 3328224, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_3328_d224_tail_div_b3_n3840_k512_d224", + "source": "tail_divisibility", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 288, + "K": 256, + "N": 8192, + "baseline_07cf_adapter_bench_iters": 1499, + "baseline_07cf_adapter_gpu_span_ms": 0.085184, + "baseline_07cf_adapter_host_enqueue_ms": 0.15552, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.00304, + "baseline_07cf_adapter_kernel_sum_ms": 0.082016, + "baseline_07cf_adapter_submission_ms": 0.15552, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.189409, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.082016 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.183808, + "synchronized_e2e_ms": 0.215872 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.085184 + }, + "host_enqueue_ms": { + "median": 0.15552 + }, + "inter_kernel_gap_ms": { + "median": 0.00304 + }, + "kernel_sum_ms": { + "median": 0.082016 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1499, + "submission_ms": { + "median": 0.15552 + }, + "synchronized_e2e_ms": { + "median": 0.189409 + } + }, + "baseline_07cf_precomputed_bench_iters": 3966, + "baseline_07cf_precomputed_gpu_span_ms": 0.041633, + "baseline_07cf_precomputed_host_enqueue_ms": 0.043776, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.041633, + "baseline_07cf_precomputed_submission_ms": 0.043776, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.092656, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.041633 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.041056, + "synchronized_e2e_ms": 0.081024 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.041633 + }, + "host_enqueue_ms": { + "median": 0.043776 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.041633 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3966, + "submission_ms": { + "median": 0.043776 + }, + "synchronized_e2e_ms": { + "median": 0.092656 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.0460692239329377, + "submission": 3.552631578947368, + "synchronized_e2e": 2.044217319979278 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.664713, + "after_init_synchronized_e2e_ms_per_call": 8.691913, + "including_init_host_enqueue_ms_per_call": 45.599151, + "including_init_synchronized_e2e_ms_per_call": 463.80323200000004, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 1.0064393, + "after_init_synchronized_e2e_ms_per_call": 1.0396594000000001, + "including_init_host_enqueue_ms_per_call": 4.699883099999999, + "including_init_synchronized_e2e_ms_per_call": 46.5507913, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.24061193, + "after_init_synchronized_e2e_ms_per_call": 0.27443404, + "including_init_host_enqueue_ms_per_call": 0.60995631, + "including_init_synchronized_e2e_ms_per_call": 4.825547230000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.164029193, + "after_init_synchronized_e2e_ms_per_call": 0.197911504, + "including_init_host_enqueue_ms_per_call": 0.200963631, + "including_init_synchronized_e2e_ms_per_call": 0.653022823, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.664713, + "synchronized_e2e_ms": 8.691913, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.389312, + "median": 0.085184, + "min": 0.082208, + "p90": 0.101632, + "sample_count": 1499 + }, + "host_enqueue_ms": { + "max": 0.895329, + "median": 0.15552, + "min": 0.134848, + "p90": 0.1864704, + "sample_count": 1499 + }, + "sample_count": 1499, + "synchronized_e2e_ms": { + "max": 0.972193, + "median": 0.189409, + "min": 0.175264, + "p90": 0.2179458, + "sample_count": 1499 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3966, + "candidate_precomputed_gpu_span_ms": 0.025632, + "candidate_precomputed_host_enqueue_ms": 0.04064, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.025632, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.04064, + "candidate_precomputed_synchronized_e2e_ms": 0.061056, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.025632 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.029824, + "synchronized_e2e_ms": 0.0456 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.025632 + }, + "host_enqueue_ms": { + "median": 0.04064 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.025632 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3966, + "submission_ms": { + "median": 0.04064 + }, + "synchronized_e2e_ms": { + "median": 0.061056 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295c2e10", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc028adb80" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.6104868913857677, + "submission": 1.1669291338582677, + "synchronized_e2e": 1.8930817610062893 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 1499, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.066912, + "candidate_public_raw_host_enqueue_ms": 0.047424, + "candidate_public_raw_inter_kernel_gap_ms": 0.000192, + "candidate_public_raw_kernel_sum_ms": 0.06672, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.047424, + "candidate_public_raw_synchronized_e2e_ms": 0.115584, + "candidate_public_raw_tflops_from_gpu_span": 72.21183357245337, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.06672 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.257824, + "synchronized_e2e_ms": 0.294048 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.066912 + }, + "host_enqueue_ms": { + "median": 0.047424 + }, + "inter_kernel_gap_ms": { + "median": 0.000192 + }, + "kernel_sum_ms": { + "median": 0.06672 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1499, + "submission_ms": { + "median": 0.047424 + }, + "synchronized_e2e_ms": { + "median": 0.115584 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.831395, + "after_init_synchronized_e2e_ms_per_call": 2.861347, + "including_init_host_enqueue_ms_per_call": 40.189417999999996, + "including_init_synchronized_e2e_ms_per_call": 40.29697, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.32582110000000003, + "after_init_synchronized_e2e_ms_per_call": 0.39016029999999996, + "including_init_host_enqueue_ms_per_call": 4.0616234, + "including_init_synchronized_e2e_ms_per_call": 4.1337226000000005, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07526371000000001, + "after_init_synchronized_e2e_ms_per_call": 0.14304163, + "including_init_host_enqueue_ms_per_call": 0.44884394, + "including_init_synchronized_e2e_ms_per_call": 0.51739786, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.050207971000000004, + "after_init_synchronized_e2e_ms_per_call": 0.118329763, + "including_init_host_enqueue_ms_per_call": 0.087565994, + "including_init_synchronized_e2e_ms_per_call": 0.155765386, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.831395, + "synchronized_e2e_ms": 2.861347, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.06768, + "median": 0.066912, + "min": 0.066016, + "p90": 0.0671752, + "sample_count": 1499 + }, + "host_enqueue_ms": { + "max": 61.723104, + "median": 0.047424, + "min": 0.039328, + "p90": 0.0639296, + "sample_count": 1499 + }, + "sample_count": 1499, + "synchronized_e2e_ms": { + "max": 61.890656, + "median": 0.115584, + "min": 0.108192, + "p90": 0.1313408, + "sample_count": 1499 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.066176, + "submission_ms": 0.066176, + "synchronized_e2e_ms": 0.104608 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.041056, + "submission_ms": 0.041056, + "synchronized_e2e_ms": 0.081024 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.183808, + "submission_ms": 0.183808, + "synchronized_e2e_ms": 0.215872 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.664713, + "submission_ms": 8.664713, + "synchronized_e2e_ms": 8.691913 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.061376, + "submission_ms": 0.061376, + "synchronized_e2e_ms": 0.077696 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.029824, + "submission_ms": 0.029824, + "synchronized_e2e_ms": 0.0456 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.218402, + "submission_ms": 1.218402, + "synchronized_e2e_ms": 1.238978 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.016129, + "submission_ms": 1.016129, + "synchronized_e2e_ms": 1.035009 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.257824, + "submission_ms": 0.257824, + "synchronized_e2e_ms": 0.294048 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.831395, + "submission_ms": 2.831395, + "synchronized_e2e_ms": 2.861347 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.344959, + "evolution_kernel_ms": 0.182432, + "evolution_speedup": 1.8909, + "evolution_tflops": 26.4857, + "expected_route": "d288_parent_splitk_hybrid_20260629_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_3328_d288_guard_overlap_b4_n8192_k256_d288", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 10930, + "measurement_schedule_sha256": "4b573812deddd01223e3ad2076a4012ca5e65c5d6518640bd27f6ff2ab27f390", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3966, + "public_pair_count": 1499, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3966, + "baseline_public_raw": 1499, + "candidate_precomputed": 3966, + "candidate_public_raw": 1499 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2188 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:adjacent_3328_d288_guard_overlap_b4_n8192_k256_d288", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.624258739076155, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.6387129706533774, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.0376997267370927, + "including_init_synchronized_e2e_speedup": 11.509630426307487, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.6646980740992876, + "including_init_synchronized_e2e_speedup": 11.261227664381735, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 1.9185606316147263, + "including_init_synchronized_e2e_speedup": 9.326569750404458, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.6725420467545429, + "including_init_synchronized_e2e_speedup": 4.1923487609756895, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.2730750836920133, + "hot_synchronized_e2e_speedup": 1.6387129706533774, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 3328288, + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_3328_d288_guard_overlap_b4_n8192_k256_d288", + "source": "guard_overlap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 3, + "D": 352, + "K": 768, + "N": 2048, + "baseline_07cf_adapter_bench_iters": 2539, + "baseline_07cf_adapter_gpu_span_ms": 0.07056, + "baseline_07cf_adapter_host_enqueue_ms": 0.1552, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.031424, + "baseline_07cf_adapter_kernel_sum_ms": 0.039072, + "baseline_07cf_adapter_submission_ms": 0.1552, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.177344, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.039072 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.213793, + "synchronized_e2e_ms": 0.233601 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.07056 + }, + "host_enqueue_ms": { + "median": 0.1552 + }, + "inter_kernel_gap_ms": { + "median": 0.031424 + }, + "kernel_sum_ms": { + "median": 0.039072 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2539, + "submission_ms": { + "median": 0.1552 + }, + "synchronized_e2e_ms": { + "median": 0.177344 + } + }, + "baseline_07cf_precomputed_bench_iters": 3085, + "baseline_07cf_precomputed_gpu_span_ms": 0.032832, + "baseline_07cf_precomputed_host_enqueue_ms": 0.045248, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.032832, + "baseline_07cf_precomputed_submission_ms": 0.045248, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.084448, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.032832 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.045504, + "synchronized_e2e_ms": 0.072608 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.032832 + }, + "host_enqueue_ms": { + "median": 0.045248 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.032832 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3085, + "submission_ms": { + "median": 0.045248 + }, + "synchronized_e2e_ms": { + "median": 0.084448 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.1491228070175437, + "submission": 3.42998585572843, + "synchronized_e2e": 2.1000378931413417 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 6.842055, + "after_init_synchronized_e2e_ms_per_call": 6.867495, + "including_init_host_enqueue_ms_per_call": 41.069962000000004, + "including_init_synchronized_e2e_ms_per_call": 41.212938, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8238855, + "after_init_synchronized_e2e_ms_per_call": 0.8463590999999999, + "including_init_host_enqueue_ms_per_call": 4.2466762000000005, + "including_init_synchronized_e2e_ms_per_call": 4.280903400000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22206855, + "after_init_synchronized_e2e_ms_per_call": 0.24424551, + "including_init_host_enqueue_ms_per_call": 0.56434762, + "including_init_synchronized_e2e_ms_per_call": 0.5876999399999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.161886855, + "after_init_synchronized_e2e_ms_per_call": 0.18403415099999998, + "including_init_host_enqueue_ms_per_call": 0.19611476200000003, + "including_init_synchronized_e2e_ms_per_call": 0.218379594, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 6.842055, + "synchronized_e2e_ms": 6.867495, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 43.712749, + "median": 0.07056, + "min": 0.060064, + "p90": 0.0922944, + "sample_count": 2539 + }, + "host_enqueue_ms": { + "max": 44.140494, + "median": 0.1552, + "min": 0.121536, + "p90": 0.23194880000000007, + "sample_count": 2539 + }, + "sample_count": 2539, + "synchronized_e2e_ms": { + "max": 44.253774, + "median": 0.177344, + "min": 0.145056, + "p90": 0.2572224, + "sample_count": 2539 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3085, + "candidate_precomputed_gpu_span_ms": 0.02848, + "candidate_precomputed_host_enqueue_ms": 0.056096, + "candidate_precomputed_inter_kernel_gap_ms": 0.003584, + "candidate_precomputed_kernel_sum_ms": 0.024832, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.056096, + "candidate_precomputed_synchronized_e2e_ms": 0.068801, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.024832 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.039904, + "synchronized_e2e_ms": 0.055584 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.02848 + }, + "host_enqueue_ms": { + "median": 0.056096 + }, + "inter_kernel_gap_ms": { + "median": 0.003584 + }, + "kernel_sum_ms": { + "median": 0.024832 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3085, + "submission_ms": { + "median": 0.056096 + }, + "synchronized_e2e_ms": { + "median": 0.068801 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282afe90", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc047356d0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.3853932584269664, + "submission": 0.8271534512264689, + "synchronized_e2e": 1.2683536576503247 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2539, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.039456, + "candidate_public_raw_host_enqueue_ms": 0.0464, + "candidate_public_raw_inter_kernel_gap_ms": 9.6e-05, + "candidate_public_raw_kernel_sum_ms": 0.039328, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.0464, + "candidate_public_raw_synchronized_e2e_ms": 0.087264, + "candidate_public_raw_tflops_from_gpu_span": 84.19223357664234, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.039328 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.076576, + "synchronized_e2e_ms": 0.10144 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.039456 + }, + "host_enqueue_ms": { + "median": 0.0464 + }, + "inter_kernel_gap_ms": { + "median": 9.6e-05 + }, + "kernel_sum_ms": { + "median": 0.039328 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2539, + "submission_ms": { + "median": 0.0464 + }, + "synchronized_e2e_ms": { + "median": 0.087264 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 40.624042, + "after_init_synchronized_e2e_ms_per_call": 40.652938, + "including_init_host_enqueue_ms_per_call": 75.201582, + "including_init_synchronized_e2e_ms_per_call": 488.74073, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 4.1041642000000005, + "after_init_synchronized_e2e_ms_per_call": 4.1438314, + "including_init_host_enqueue_ms_per_call": 7.561918199999999, + "including_init_synchronized_e2e_ms_per_call": 48.9526106, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.45217642, + "after_init_synchronized_e2e_ms_per_call": 0.49292073999999997, + "including_init_host_enqueue_ms_per_call": 0.79795182, + "including_init_synchronized_e2e_ms_per_call": 4.97379866, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.08697764200000001, + "after_init_synchronized_e2e_ms_per_call": 0.12782967399999998, + "including_init_host_enqueue_ms_per_call": 0.121555182, + "including_init_synchronized_e2e_ms_per_call": 0.575917466, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 40.624042, + "synchronized_e2e_ms": 40.652938, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.040064, + "median": 0.039456, + "min": 0.038976, + "p90": 0.0396482, + "sample_count": 2539 + }, + "host_enqueue_ms": { + "max": 34.294468, + "median": 0.0464, + "min": 0.034208, + "p90": 0.07127040000000005, + "sample_count": 2539 + }, + "sample_count": 2539, + "synchronized_e2e_ms": { + "max": 34.444868, + "median": 0.087264, + "min": 0.0776, + "p90": 0.10738560000000004, + "sample_count": 2539 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.067329, + "submission_ms": 0.067329, + "synchronized_e2e_ms": 0.092417 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.045504, + "submission_ms": 0.045504, + "synchronized_e2e_ms": 0.072608 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.213793, + "submission_ms": 0.213793, + "synchronized_e2e_ms": 0.233601 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 6.842055, + "submission_ms": 6.842055, + "synchronized_e2e_ms": 6.867495 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.06832, + "submission_ms": 0.06832, + "synchronized_e2e_ms": 0.084896 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.039904, + "submission_ms": 0.039904, + "synchronized_e2e_ms": 0.055584 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.325889, + "submission_ms": 1.325889, + "synchronized_e2e_ms": 1.348961 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.109057, + "submission_ms": 1.109057, + "synchronized_e2e_ms": 1.131265 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.076576, + "submission_ms": 0.076576, + "synchronized_e2e_ms": 0.10144 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 40.624042, + "submission_ms": 40.624042, + "synchronized_e2e_ms": 40.652938 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.582335, + "evolution_kernel_ms": 0.183872, + "evolution_speedup": 3.1671, + "evolution_tflops": 18.0663, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_3328_d352_random_legal_b3_n2048_k768_d352", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 11248, + "measurement_schedule_sha256": "69c51af4623e39ab936bd0d35058c61f08e520d29364b493e03380ae0fa35aea", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3085, + "public_pair_count": 2539, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3085, + "baseline_public_raw": 2539, + "candidate_precomputed": 3085, + "candidate_public_raw": 2539 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2250 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:adjacent_3328_d352_random_legal_b3_n2048_k768_d352", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.152808988764045, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.0322698936560326, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.16892985692694584, + "including_init_synchronized_e2e_speedup": 0.08432474616961022, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.20424554435298695, + "including_init_synchronized_e2e_speedup": 0.08744995103488926, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.4955066609694695, + "including_init_synchronized_e2e_speedup": 0.11815917373704064, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.439682549765401, + "including_init_synchronized_e2e_speedup": 0.3791855723993618, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.7883211678832116, + "hot_synchronized_e2e_speedup": 2.0322698936560326, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 3328352, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_3328_d352_random_legal_b3_n2048_k768_d352", + "source": "random_legal", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 416, + "K": 8192, + "N": 384, + "baseline_07cf_adapter_bench_iters": 2307, + "baseline_07cf_adapter_gpu_span_ms": 0.258528, + "baseline_07cf_adapter_host_enqueue_ms": 0.152672, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.030528, + "baseline_07cf_adapter_kernel_sum_ms": 0.227969, + "baseline_07cf_adapter_submission_ms": 0.152672, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.361313, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.227969 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.219136, + "synchronized_e2e_ms": 0.418912 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.258528 + }, + "host_enqueue_ms": { + "median": 0.152672 + }, + "inter_kernel_gap_ms": { + "median": 0.030528 + }, + "kernel_sum_ms": { + "median": 0.227969 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2307, + "submission_ms": { + "median": 0.152672 + }, + "synchronized_e2e_ms": { + "median": 0.361313 + } + }, + "baseline_07cf_precomputed_bench_iters": 2447, + "baseline_07cf_precomputed_gpu_span_ms": 0.255968, + "baseline_07cf_precomputed_host_enqueue_ms": 0.045952, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.255968, + "baseline_07cf_precomputed_submission_ms": 0.045952, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.307936, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.255968 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.053184, + "synchronized_e2e_ms": 0.248288 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.255968 + }, + "host_enqueue_ms": { + "median": 0.045952 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.255968 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2447, + "submission_ms": { + "median": 0.045952 + }, + "synchronized_e2e_ms": { + "median": 0.307936 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.0100012501562696, + "submission": 3.322423398328691, + "synchronized_e2e": 1.1733379663306662 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 6.452679, + "after_init_synchronized_e2e_ms_per_call": 6.636327, + "including_init_host_enqueue_ms_per_call": 40.902123, + "including_init_synchronized_e2e_ms_per_call": 490.677821, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.7826727, + "after_init_synchronized_e2e_ms_per_call": 0.9888144000000001, + "including_init_host_enqueue_ms_per_call": 4.227617100000001, + "including_init_synchronized_e2e_ms_per_call": 49.392963800000004, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.21567207, + "after_init_synchronized_e2e_ms_per_call": 0.42406314, + "including_init_host_enqueue_ms_per_call": 0.5601665100000001, + "including_init_synchronized_e2e_ms_per_call": 5.26447808, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.158972007, + "after_init_synchronized_e2e_ms_per_call": 0.36758801399999996, + "including_init_host_enqueue_ms_per_call": 0.193421451, + "including_init_synchronized_e2e_ms_per_call": 0.851629508, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 6.452679, + "synchronized_e2e_ms": 6.636327, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 35.993989, + "median": 0.258528, + "min": 0.250048, + "p90": 0.275008, + "sample_count": 2307 + }, + "host_enqueue_ms": { + "max": 36.087462, + "median": 0.152672, + "min": 0.130464, + "p90": 0.19122660000000002, + "sample_count": 2307 + }, + "sample_count": 2307, + "synchronized_e2e_ms": { + "max": 36.27143, + "median": 0.361313, + "min": 0.341696, + "p90": 0.3963332, + "sample_count": 2307 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 4, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 2447, + "candidate_precomputed_gpu_span_ms": 0.03312, + "candidate_precomputed_host_enqueue_ms": 0.064608, + "candidate_precomputed_inter_kernel_gap_ms": 0.005472, + "candidate_precomputed_kernel_sum_ms": 0.02736, + "candidate_precomputed_launch_count": 3, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.064608, + "candidate_precomputed_synchronized_e2e_ms": 0.077152, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.02736 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.062305, + "synchronized_e2e_ms": 0.077505 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 3.0 + }, + "gpu_span_ms": { + "median": 0.03312 + }, + "host_enqueue_ms": { + "median": 0.064608 + }, + "inter_kernel_gap_ms": { + "median": 0.005472 + }, + "kernel_sum_ms": { + "median": 0.02736 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2447, + "submission_ms": { + "median": 0.064608 + }, + "synchronized_e2e_ms": { + "median": 0.077152 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295bfe00", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295bef30" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.317874396135266, + "submission": 0.7845468053491828, + "synchronized_e2e": 1.226060244711738 + }, + "candidate_public_raw_assignment_launch_count": 3, + "candidate_public_raw_bench_iters": 2307, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.043648, + "candidate_public_raw_host_enqueue_ms": 0.050688, + "candidate_public_raw_inter_kernel_gap_ms": 9.6e-05, + "candidate_public_raw_kernel_sum_ms": 0.043552, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.050688, + "candidate_public_raw_synchronized_e2e_ms": 0.094593, + "candidate_public_raw_tflops_from_gpu_span": 59.962557184750736, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.043552 + }, + "activity_count": { + "median": 4.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.089024, + "synchronized_e2e_ms": 0.112992 + }, + "correlated_kernel_activity_count": { + "median": 4.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.043648 + }, + "host_enqueue_ms": { + "median": 0.050688 + }, + "inter_kernel_gap_ms": { + "median": 9.6e-05 + }, + "kernel_sum_ms": { + "median": 0.043552 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2307, + "submission_ms": { + "median": 0.050688 + }, + "synchronized_e2e_ms": { + "median": 0.094593 + } + }, + "candidate_public_raw_total_launch_count": 4, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 76.705488, + "after_init_synchronized_e2e_ms_per_call": 76.74536, + "including_init_host_enqueue_ms_per_call": 111.579572, + "including_init_synchronized_e2e_ms_per_call": 111.688436, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 7.716168000000001, + "after_init_synchronized_e2e_ms_per_call": 7.759669700000001, + "including_init_host_enqueue_ms_per_call": 11.2035764, + "including_init_synchronized_e2e_ms_per_call": 11.253977299999999, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8172360000000001, + "after_init_synchronized_e2e_ms_per_call": 0.86110067, + "including_init_host_enqueue_ms_per_call": 1.16597684, + "including_init_synchronized_e2e_ms_per_call": 1.2105314299999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.1273428, + "after_init_synchronized_e2e_ms_per_call": 0.171243767, + "including_init_host_enqueue_ms_per_call": 0.162216884, + "including_init_synchronized_e2e_ms_per_call": 0.206186843, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 76.705488, + "synchronized_e2e_ms": 76.74536, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.044576, + "median": 0.043648, + "min": 0.043008, + "p90": 0.043968, + "sample_count": 2307 + }, + "host_enqueue_ms": { + "max": 0.425057, + "median": 0.050688, + "min": 0.038976, + "p90": 0.06842880000000001, + "sample_count": 2307 + }, + "sample_count": 2307, + "synchronized_e2e_ms": { + "max": 0.504033, + "median": 0.094593, + "min": 0.084064, + "p90": 0.11056, + "sample_count": 2307 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.079104, + "submission_ms": 0.079104, + "synchronized_e2e_ms": 0.27328 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.053184, + "submission_ms": 0.053184, + "synchronized_e2e_ms": 0.248288 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.219136, + "submission_ms": 0.219136, + "synchronized_e2e_ms": 0.418912 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 6.452679, + "submission_ms": 6.452679, + "synchronized_e2e_ms": 6.636327 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.099328, + "submission_ms": 0.099328, + "synchronized_e2e_ms": 0.116896 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.062305, + "submission_ms": 0.062305, + "synchronized_e2e_ms": 0.077505 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.468514, + "submission_ms": 1.468514, + "synchronized_e2e_ms": 1.49501 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.364929, + "submission_ms": 1.364929, + "synchronized_e2e_ms": 1.391073 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.089024, + "submission_ms": 0.089024, + "synchronized_e2e_ms": 0.112992 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 76.705488, + "submission_ms": 76.705488, + "synchronized_e2e_ms": 76.74536 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 41.538691, + "evolution_kernel_ms": 0.276031, + "evolution_speedup": 150.4856, + "evolution_tflops": 9.4817, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_3328_d416_highk_low_n_b1_n384_k8192_d416", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 9508, + "measurement_schedule_sha256": "b5dfa2e233ecbcf50083c924f1b1926bff1b2e07cec266080299d208f123079b", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 2447, + "public_pair_count": 2307, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 2447, + "baseline_public_raw": 2307, + "candidate_precomputed": 2447, + "candidate_public_raw": 2307 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 1904 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:adjacent_3328_d416_highk_low_n_b1_n384_k8192_d416", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 7.728502415458937, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 3.819658959965325, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.08647202905817367, + "including_init_synchronized_e2e_speedup": 4.393273274952118, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.12742996006647037, + "including_init_synchronized_e2e_speedup": 4.388934017131882, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.4924663918795929, + "including_init_synchronized_e2e_speedup": 4.348898301632698, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.1465774809777454, + "including_init_synchronized_e2e_speedup": 4.1303775527519955, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 5.923020527859237, + "hot_synchronized_e2e_speedup": 3.819658959965325, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 3328416, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_3328_d416_highk_low_n_b1_n384_k8192_d416", + "source": "request_specific", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 480, + "K": 256, + "N": 128, + "baseline_07cf_adapter_bench_iters": 5527, + "baseline_07cf_adapter_gpu_span_ms": 0.05568, + "baseline_07cf_adapter_host_enqueue_ms": 0.152128, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.04176, + "baseline_07cf_adapter_kernel_sum_ms": 0.01392, + "baseline_07cf_adapter_submission_ms": 0.152128, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.170176, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.01392 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.340544, + "synchronized_e2e_ms": 0.369888 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.05568 + }, + "host_enqueue_ms": { + "median": 0.152128 + }, + "inter_kernel_gap_ms": { + "median": 0.04176 + }, + "kernel_sum_ms": { + "median": 0.01392 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 5527, + "submission_ms": { + "median": 0.152128 + }, + "synchronized_e2e_ms": { + "median": 0.170176 + } + }, + "baseline_07cf_precomputed_bench_iters": 7842, + "baseline_07cf_precomputed_gpu_span_ms": 0.012896, + "baseline_07cf_precomputed_host_enqueue_ms": 0.0433285, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.012896, + "baseline_07cf_precomputed_submission_ms": 0.0433285, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.061568, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.012896 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.066016, + "synchronized_e2e_ms": 0.087296 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.012896 + }, + "host_enqueue_ms": { + "median": 0.0433285 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.012896 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 7842, + "submission_ms": { + "median": 0.0433285 + }, + "synchronized_e2e_ms": { + "median": 0.061568 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 4.317617866004963, + "submission": 3.5110377695973787, + "synchronized_e2e": 2.764033264033264 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.132232, + "after_init_synchronized_e2e_ms_per_call": 8.173288, + "including_init_host_enqueue_ms_per_call": 43.6835, + "including_init_synchronized_e2e_ms_per_call": 43.804525, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9501383999999999, + "after_init_synchronized_e2e_ms_per_call": 0.9704872, + "including_init_host_enqueue_ms_per_call": 4.5052652, + "including_init_synchronized_e2e_ms_per_call": 4.5336109, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23192904000000003, + "after_init_synchronized_e2e_ms_per_call": 0.25020712, + "including_init_host_enqueue_ms_per_call": 0.5874417200000001, + "including_init_synchronized_e2e_ms_per_call": 0.60651949, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.160108104, + "after_init_synchronized_e2e_ms_per_call": 0.17817911199999997, + "including_init_host_enqueue_ms_per_call": 0.19565937200000003, + "including_init_synchronized_e2e_ms_per_call": 0.21381034899999998, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.132232, + "synchronized_e2e_ms": 8.173288, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 9.996842, + "median": 0.05568, + "min": 0.045632, + "p90": 0.08849320000000002, + "sample_count": 5527 + }, + "host_enqueue_ms": { + "max": 35.671301, + "median": 0.152128, + "min": 0.122016, + "p90": 0.2761664000000006, + "sample_count": 5527 + }, + "sample_count": 5527, + "synchronized_e2e_ms": { + "max": 35.703013, + "median": 0.170176, + "min": 0.136832, + "p90": 0.306765800000001, + "sample_count": 5527 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 7842, + "candidate_precomputed_gpu_span_ms": 0.01984, + "candidate_precomputed_host_enqueue_ms": 0.053984, + "candidate_precomputed_inter_kernel_gap_ms": 0.00512, + "candidate_precomputed_kernel_sum_ms": 0.014688, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.053984, + "candidate_precomputed_synchronized_e2e_ms": 0.065632, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.014688 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.065376, + "synchronized_e2e_ms": 0.084704 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.01984 + }, + "host_enqueue_ms": { + "median": 0.053984 + }, + "inter_kernel_gap_ms": { + "median": 0.00512 + }, + "kernel_sum_ms": { + "median": 0.014688 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 7842, + "submission_ms": { + "median": 0.053984 + }, + "synchronized_e2e_ms": { + "median": 0.065632 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282c66c0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282c6d50" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 0.9225806451612903, + "submission": 0.8393598103141672, + "synchronized_e2e": 0.976596782057533 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 5527, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.018304, + "candidate_public_raw_host_enqueue_ms": 0.045312, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.01824, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.045312, + "candidate_public_raw_synchronized_e2e_ms": 0.064096, + "candidate_public_raw_tflops_from_gpu_span": 1.7186013986013984, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.01824 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.11232, + "synchronized_e2e_ms": 0.140929 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.018304 + }, + "host_enqueue_ms": { + "median": 0.045312 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.01824 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 5527, + "submission_ms": { + "median": 0.045312 + }, + "synchronized_e2e_ms": { + "median": 0.064096 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.889923, + "after_init_synchronized_e2e_ms_per_call": 2.927075, + "including_init_host_enqueue_ms_per_call": 38.727943, + "including_init_synchronized_e2e_ms_per_call": 487.027288, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.32977310000000004, + "after_init_synchronized_e2e_ms_per_call": 0.3503939, + "including_init_host_enqueue_ms_per_call": 3.9135751000000005, + "including_init_synchronized_e2e_ms_per_call": 48.7604152, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07375811, + "after_init_synchronized_e2e_ms_per_call": 0.09272579, + "including_init_host_enqueue_ms_per_call": 0.4321383100000001, + "including_init_synchronized_e2e_ms_per_call": 4.93372792, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.048156610999999995, + "after_init_synchronized_e2e_ms_per_call": 0.066958979, + "including_init_host_enqueue_ms_per_call": 0.083994631, + "including_init_synchronized_e2e_ms_per_call": 0.551059192, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.889923, + "synchronized_e2e_ms": 2.927075, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.019072, + "median": 0.018304, + "min": 0.0176, + "p90": 0.018496, + "sample_count": 5527 + }, + "host_enqueue_ms": { + "max": 25.611482, + "median": 0.045312, + "min": 0.033664, + "p90": 0.0867008000000001, + "sample_count": 5527 + }, + "sample_count": 5527, + "synchronized_e2e_ms": { + "max": 35.842661, + "median": 0.064096, + "min": 0.054528, + "p90": 0.11165440000000015, + "sample_count": 5527 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.102497, + "submission_ms": 0.102497, + "synchronized_e2e_ms": 0.128801 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.066016, + "submission_ms": 0.066016, + "synchronized_e2e_ms": 0.087296 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.340544, + "submission_ms": 0.340544, + "synchronized_e2e_ms": 0.369888 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.132232, + "submission_ms": 8.132232, + "synchronized_e2e_ms": 8.173288 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.104224, + "submission_ms": 0.104224, + "synchronized_e2e_ms": 0.127616 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.065376, + "submission_ms": 0.065376, + "synchronized_e2e_ms": 0.084704 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.644002, + "submission_ms": 1.644002, + "synchronized_e2e_ms": 1.680162 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.494049, + "submission_ms": 1.494049, + "synchronized_e2e_ms": 1.531041 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.11232, + "submission_ms": 0.11232, + "synchronized_e2e_ms": 0.140929 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.889923, + "submission_ms": 2.889923, + "synchronized_e2e_ms": 2.927075 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.488575, + "evolution_kernel_ms": 0.177056, + "evolution_speedup": 2.7594, + "evolution_tflops": 0.1777, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_3328_d480_min_boundary_b1_n128_k256_d480", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 26738, + "measurement_schedule_sha256": "b2eb17b035664122487752e4a7890526b5d2f36f0629b82a6464b1e164ce3e59", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 7842, + "public_pair_count": 5527, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 7842, + "baseline_public_raw": 5527, + "candidate_precomputed": 7842, + "candidate_public_raw": 5527 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 5350 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:adjacent_3328_d480_min_boundary_b1_n128_k256_d480", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.65, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.6550174737893157, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.79230562933987, + "including_init_synchronized_e2e_speedup": 0.08994265019499277, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.769703468011287, + "including_init_synchronized_e2e_speedup": 0.09297728252322184, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.69835522566052, + "including_init_synchronized_e2e_speedup": 0.12293330719380247, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.661018950124672, + "including_init_synchronized_e2e_speedup": 0.3879988794379824, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 3.0419580419580416, + "hot_synchronized_e2e_speedup": 2.6550174737893157, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 3328480, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_3328_d480_min_boundary_b1_n128_k256_d480", + "source": "guard_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 48, + "K": 256, + "N": 256, + "baseline_07cf_adapter_bench_iters": 10666, + "baseline_07cf_adapter_gpu_span_ms": 0.05264, + "baseline_07cf_adapter_host_enqueue_ms": 0.154753, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.044864, + "baseline_07cf_adapter_kernel_sum_ms": 0.007808, + "baseline_07cf_adapter_submission_ms": 0.154753, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.174112, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.007808 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.242848, + "synchronized_e2e_ms": 0.28576 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.05264 + }, + "host_enqueue_ms": { + "median": 0.154753 + }, + "inter_kernel_gap_ms": { + "median": 0.044864 + }, + "kernel_sum_ms": { + "median": 0.007808 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 10666, + "submission_ms": { + "median": 0.154753 + }, + "synchronized_e2e_ms": { + "median": 0.174112 + } + }, + "baseline_07cf_precomputed_bench_iters": 14917, + "baseline_07cf_precomputed_gpu_span_ms": 0.00672, + "baseline_07cf_precomputed_host_enqueue_ms": 0.045152, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.00672, + "baseline_07cf_precomputed_submission_ms": 0.045152, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.059616, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.00672 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.052832, + "synchronized_e2e_ms": 0.072512 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.00672 + }, + "host_enqueue_ms": { + "median": 0.045152 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.00672 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 14917, + "submission_ms": { + "median": 0.045152 + }, + "synchronized_e2e_ms": { + "median": 0.059616 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 7.833333333333333, + "submission": 3.427378632175762, + "synchronized_e2e": 2.9205582393988188 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.440104, + "after_init_synchronized_e2e_ms_per_call": 7.471016, + "including_init_host_enqueue_ms_per_call": 44.374542, + "including_init_synchronized_e2e_ms_per_call": 462.58233500000006, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8832881, + "after_init_synchronized_e2e_ms_per_call": 0.9038024, + "including_init_host_enqueue_ms_per_call": 4.5767319, + "including_init_synchronized_e2e_ms_per_call": 46.414934300000006, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22760650999999998, + "after_init_synchronized_e2e_ms_per_call": 0.24708103999999997, + "including_init_host_enqueue_ms_per_call": 0.59695089, + "including_init_synchronized_e2e_ms_per_call": 4.79819423, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.162038351, + "after_init_synchronized_e2e_ms_per_call": 0.18140890399999998, + "including_init_host_enqueue_ms_per_call": 0.198972789, + "including_init_synchronized_e2e_ms_per_call": 0.636520223, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.440104, + "synchronized_e2e_ms": 7.471016, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 61.498017, + "median": 0.05264, + "min": 0.041344, + "p90": 0.076736, + "sample_count": 10666 + }, + "host_enqueue_ms": { + "max": 111.429491, + "median": 0.154753, + "min": 0.1224, + "p90": 0.2149125, + "sample_count": 10666 + }, + "sample_count": 10666, + "synchronized_e2e_ms": { + "max": 111.686132, + "median": 0.174112, + "min": 0.140992, + "p90": 0.2391685, + "sample_count": 10666 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 14917, + "candidate_precomputed_gpu_span_ms": 0.014304, + "candidate_precomputed_host_enqueue_ms": 0.057792, + "candidate_precomputed_inter_kernel_gap_ms": 0.00704, + "candidate_precomputed_kernel_sum_ms": 0.007296, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.057792, + "candidate_precomputed_synchronized_e2e_ms": 0.070176, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.007296 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.059168, + "synchronized_e2e_ms": 0.098752 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.014304 + }, + "host_enqueue_ms": { + "median": 0.057792 + }, + "inter_kernel_gap_ms": { + "median": 0.00704 + }, + "kernel_sum_ms": { + "median": 0.007296 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 14917, + "submission_ms": { + "median": 0.057792 + }, + "synchronized_e2e_ms": { + "median": 0.070176 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295c3d10", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282c69f0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 0.6845637583892618, + "submission": 0.8388704318936877, + "synchronized_e2e": 0.9042407660738714 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 10666, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.009792, + "candidate_public_raw_host_enqueue_ms": 0.04848, + "candidate_public_raw_inter_kernel_gap_ms": 0.000128, + "candidate_public_raw_kernel_sum_ms": 0.009664, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.04848, + "candidate_public_raw_synchronized_e2e_ms": 0.063456, + "candidate_public_raw_tflops_from_gpu_span": 0.6425098039215686, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.009664 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.093344, + "synchronized_e2e_ms": 0.118176 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.009792 + }, + "host_enqueue_ms": { + "median": 0.04848 + }, + "inter_kernel_gap_ms": { + "median": 0.000128 + }, + "kernel_sum_ms": { + "median": 0.009664 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 10666, + "submission_ms": { + "median": 0.04848 + }, + "synchronized_e2e_ms": { + "median": 0.063456 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.566787, + "after_init_synchronized_e2e_ms_per_call": 3.595107, + "including_init_host_enqueue_ms_per_call": 40.924809999999994, + "including_init_synchronized_e2e_ms_per_call": 41.03073, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.4003107, + "after_init_synchronized_e2e_ms_per_call": 0.41662109999999997, + "including_init_host_enqueue_ms_per_call": 4.136113, + "including_init_synchronized_e2e_ms_per_call": 4.160183399999999, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.08366307, + "after_init_synchronized_e2e_ms_per_call": 0.09877251, + "including_init_host_enqueue_ms_per_call": 0.45724329999999996, + "including_init_synchronized_e2e_ms_per_call": 0.47312874, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.051998306999999994, + "after_init_synchronized_e2e_ms_per_call": 0.066987651, + "including_init_host_enqueue_ms_per_call": 0.08935632999999998, + "including_init_synchronized_e2e_ms_per_call": 0.104423274, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.566787, + "synchronized_e2e_ms": 3.595107, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.010432, + "median": 0.009792, + "min": 0.009184, + "p90": 0.009952, + "sample_count": 10666 + }, + "host_enqueue_ms": { + "max": 118.834075, + "median": 0.04848, + "min": 0.037312, + "p90": 0.07460800000000001, + "sample_count": 10666 + }, + "sample_count": 10666, + "synchronized_e2e_ms": { + "max": 119.078683, + "median": 0.063456, + "min": 0.050208, + "p90": 0.09504, + "sample_count": 10666 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.078753, + "submission_ms": 0.078753, + "synchronized_e2e_ms": 0.101057 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.052832, + "submission_ms": 0.052832, + "synchronized_e2e_ms": 0.072512 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.242848, + "submission_ms": 0.242848, + "synchronized_e2e_ms": 0.28576 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.440104, + "submission_ms": 7.440104, + "synchronized_e2e_ms": 7.471016 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.105504, + "submission_ms": 0.105504, + "synchronized_e2e_ms": 0.128928 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.059168, + "submission_ms": 0.059168, + "synchronized_e2e_ms": 0.098752 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.483778, + "submission_ms": 1.483778, + "synchronized_e2e_ms": 1.510914 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.223841, + "submission_ms": 1.223841, + "synchronized_e2e_ms": 1.248609 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.093344, + "submission_ms": 0.093344, + "synchronized_e2e_ms": 0.118176 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.566787, + "submission_ms": 3.566787, + "synchronized_e2e_ms": 3.595107 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.434063, + "evolution_kernel_ms": 0.168656, + "evolution_speedup": 2.5737, + "evolution_tflops": 0.0373, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_3328_d48_small_boundary_b1_n256_k256_d48", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 51166, + "measurement_schedule_sha256": "0b0d396b6f3cfd10e9667faede8ce10e17d610c9c16323cc56c67622c0610a35", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 14917, + "public_pair_count": 10666, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 14917, + "baseline_public_raw": 10666, + "candidate_precomputed": 14917, + "candidate_public_raw": 10666 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 10236 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:adjacent_3328_d48_small_boundary_b1_n256_k256_d48", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.4697986577181208, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.743822491174987, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.0781067155998416, + "including_init_synchronized_e2e_speedup": 11.274045940688847, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.1693630015378482, + "including_init_synchronized_e2e_speedup": 11.15694425875552, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.5015162619639817, + "including_init_synchronized_e2e_speedup": 10.14141358227361, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.7080947203238996, + "including_init_synchronized_e2e_speedup": 6.095578108382237, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 5.375816993464052, + "hot_synchronized_e2e_speedup": 2.743822491174987, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 3328481, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_3328_d48_small_boundary_b1_n256_k256_d48", + "source": "guard_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 64, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 5, + "D": 112, + "K": 512, + "N": 2176, + "baseline_07cf_adapter_bench_iters": 4883, + "baseline_07cf_adapter_gpu_span_ms": 0.056704, + "baseline_07cf_adapter_host_enqueue_ms": 0.14704, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.034272, + "baseline_07cf_adapter_kernel_sum_ms": 0.022432, + "baseline_07cf_adapter_submission_ms": 0.14704, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.166144, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.022432 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.326273, + "synchronized_e2e_ms": 0.356545 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.056704 + }, + "host_enqueue_ms": { + "median": 0.14704 + }, + "inter_kernel_gap_ms": { + "median": 0.034272 + }, + "kernel_sum_ms": { + "median": 0.022432 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 4883, + "submission_ms": { + "median": 0.14704 + }, + "synchronized_e2e_ms": { + "median": 0.166144 + } + }, + "baseline_07cf_precomputed_bench_iters": 7111, + "baseline_07cf_precomputed_gpu_span_ms": 0.014048, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042208, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.014048, + "baseline_07cf_precomputed_submission_ms": 0.042208, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.06288, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.014048 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.051296, + "synchronized_e2e_ms": 0.069056 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.014048 + }, + "host_enqueue_ms": { + "median": 0.042208 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.014048 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 7111, + "submission_ms": { + "median": 0.042208 + }, + "synchronized_e2e_ms": { + "median": 0.06288 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 4.0364464692482915, + "submission": 3.483699772554966, + "synchronized_e2e": 2.6422391857506358 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.045576, + "after_init_synchronized_e2e_ms_per_call": 8.083496, + "including_init_host_enqueue_ms_per_call": 42.273483, + "including_init_synchronized_e2e_ms_per_call": 42.428939, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9368936000000001, + "after_init_synchronized_e2e_ms_per_call": 0.9578792, + "including_init_host_enqueue_ms_per_call": 4.3596843, + "including_init_synchronized_e2e_ms_per_call": 4.3924235000000005, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22602536, + "after_init_synchronized_e2e_ms_per_call": 0.24531751999999998, + "including_init_host_enqueue_ms_per_call": 0.56830443, + "including_init_synchronized_e2e_ms_per_call": 0.58877195, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15493853600000002, + "after_init_synchronized_e2e_ms_per_call": 0.17406135199999997, + "including_init_host_enqueue_ms_per_call": 0.18916644300000002, + "including_init_synchronized_e2e_ms_per_call": 0.20840679499999998, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.045576, + "synchronized_e2e_ms": 8.083496, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 56.575354, + "median": 0.056704, + "min": 0.04688, + "p90": 0.08195840000000004, + "sample_count": 4883 + }, + "host_enqueue_ms": { + "max": 57.157916, + "median": 0.14704, + "min": 0.120704, + "p90": 0.23811860000000007, + "sample_count": 4883 + }, + "sample_count": 4883, + "synchronized_e2e_ms": { + "max": 57.29206, + "median": 0.166144, + "min": 0.136832, + "p90": 0.26398160000000004, + "sample_count": 4883 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 7111, + "candidate_precomputed_gpu_span_ms": 0.014848, + "candidate_precomputed_host_enqueue_ms": 0.05392, + "candidate_precomputed_inter_kernel_gap_ms": 0.004224, + "candidate_precomputed_kernel_sum_ms": 0.010656, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.05392, + "candidate_precomputed_synchronized_e2e_ms": 0.066016, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.010656 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.048736, + "synchronized_e2e_ms": 0.064768 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.014848 + }, + "host_enqueue_ms": { + "median": 0.05392 + }, + "inter_kernel_gap_ms": { + "median": 0.004224 + }, + "kernel_sum_ms": { + "median": 0.010656 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 7111, + "submission_ms": { + "median": 0.05392 + }, + "synchronized_e2e_ms": { + "median": 0.066016 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffe4324f470", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282c5bb0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.4073275862068966, + "submission": 0.8243323442136499, + "synchronized_e2e": 1.0159961221522056 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 4883, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.020896, + "candidate_public_raw_host_enqueue_ms": 0.044448, + "candidate_public_raw_inter_kernel_gap_ms": 9.6e-05, + "candidate_public_raw_kernel_sum_ms": 0.0208, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.044448, + "candidate_public_raw_synchronized_e2e_ms": 0.067072, + "candidate_public_raw_tflops_from_gpu_span": 59.7150382848392, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.0208 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.101408, + "synchronized_e2e_ms": 0.131808 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.020896 + }, + "host_enqueue_ms": { + "median": 0.044448 + }, + "inter_kernel_gap_ms": { + "median": 9.6e-05 + }, + "kernel_sum_ms": { + "median": 0.0208 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 4883, + "submission_ms": { + "median": 0.044448 + }, + "synchronized_e2e_ms": { + "median": 0.067072 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.952355, + "after_init_synchronized_e2e_ms_per_call": 2.985571, + "including_init_host_enqueue_ms_per_call": 37.529894999999996, + "including_init_synchronized_e2e_ms_per_call": 451.073363, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.3352387, + "after_init_synchronized_e2e_ms_per_call": 0.3589219, + "including_init_host_enqueue_ms_per_call": 3.7929927, + "including_init_synchronized_e2e_ms_per_call": 45.1677011, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07352707, + "after_init_synchronized_e2e_ms_per_call": 0.09625699000000001, + "including_init_host_enqueue_ms_per_call": 0.4193024699999999, + "including_init_synchronized_e2e_ms_per_call": 4.57713491, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.047355907, + "after_init_synchronized_e2e_ms_per_call": 0.069990499, + "including_init_host_enqueue_ms_per_call": 0.081933447, + "including_init_synchronized_e2e_ms_per_call": 0.518078291, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.952355, + "synchronized_e2e_ms": 2.985571, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.021568, + "median": 0.020896, + "min": 0.02016, + "p90": 0.021088, + "sample_count": 4883 + }, + "host_enqueue_ms": { + "max": 45.831888, + "median": 0.044448, + "min": 0.034848, + "p90": 0.07391360000000001, + "sample_count": 4883 + }, + "sample_count": 4883, + "synchronized_e2e_ms": { + "max": 45.997552, + "median": 0.067072, + "min": 0.058976, + "p90": 0.094816, + "sample_count": 4883 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.087232, + "submission_ms": 0.087232, + "synchronized_e2e_ms": 0.105088 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.051296, + "submission_ms": 0.051296, + "synchronized_e2e_ms": 0.069056 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.326273, + "submission_ms": 0.326273, + "synchronized_e2e_ms": 0.356545 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.045576, + "submission_ms": 8.045576, + "synchronized_e2e_ms": 8.083496 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.085408, + "submission_ms": 0.085408, + "synchronized_e2e_ms": 0.103008 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.048736, + "submission_ms": 0.048736, + "synchronized_e2e_ms": 0.064768 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.594209, + "submission_ms": 1.594209, + "synchronized_e2e_ms": 1.629409 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.47357, + "submission_ms": 1.47357, + "synchronized_e2e_ms": 1.499778 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.101408, + "submission_ms": 0.101408, + "synchronized_e2e_ms": 0.131808 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.952355, + "submission_ms": 2.952355, + "synchronized_e2e_ms": 2.985571 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.338463, + "evolution_kernel_ms": 0.171519, + "evolution_speedup": 1.9733, + "evolution_tflops": 7.275, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_5600_d112_odd_tile_tail_b5_n2176_k512_d112", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 23988, + "measurement_schedule_sha256": "4005c23667e6648bb598514b1c6525a0402c76a799903e5b540352bb95ef520d", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 7111, + "public_pair_count": 4883, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 7111, + "baseline_public_raw": 4883, + "candidate_precomputed": 7111, + "candidate_public_raw": 4883 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 4800 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:adjacent_5600_d112_odd_tile_tail_b5_n2176_k512_d112", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.9461206896551724, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.477099236641221, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.707520939880512, + "including_init_synchronized_e2e_speedup": 0.09406216921747163, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.6687677737134456, + "including_init_synchronized_e2e_speedup": 0.09724700157476025, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.5485683689049488, + "including_init_synchronized_e2e_speedup": 0.12863329606336643, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.4869282900812006, + "including_init_synchronized_e2e_speedup": 0.4022689207797745, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.713629402756508, + "hot_synchronized_e2e_speedup": 2.477099236641221, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 5601121, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_5600_d112_odd_tile_tail_b5_n2176_k512_d112", + "source": "tail_divisibility", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 8, + "D": 128, + "K": 256, + "N": 8320, + "baseline_07cf_adapter_bench_iters": 1779, + "baseline_07cf_adapter_gpu_span_ms": 0.06784, + "baseline_07cf_adapter_host_enqueue_ms": 0.164705, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.010432, + "baseline_07cf_adapter_kernel_sum_ms": 0.057312, + "baseline_07cf_adapter_submission_ms": 0.164705, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.18496, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.057312 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.24592, + "synchronized_e2e_ms": 0.269152 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.06784 + }, + "host_enqueue_ms": { + "median": 0.164705 + }, + "inter_kernel_gap_ms": { + "median": 0.010432 + }, + "kernel_sum_ms": { + "median": 0.057312 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1779, + "submission_ms": { + "median": 0.164705 + }, + "synchronized_e2e_ms": { + "median": 0.18496 + } + }, + "baseline_07cf_precomputed_bench_iters": 5546, + "baseline_07cf_precomputed_gpu_span_ms": 0.018464, + "baseline_07cf_precomputed_host_enqueue_ms": 0.045967999999999995, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.018464, + "baseline_07cf_precomputed_submission_ms": 0.045967999999999995, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.0696645, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.018464 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.065696, + "synchronized_e2e_ms": 0.08768 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.018464 + }, + "host_enqueue_ms": { + "median": 0.045967999999999995 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.018464 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5546, + "submission_ms": { + "median": 0.045967999999999995 + }, + "synchronized_e2e_ms": { + "median": 0.0696645 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 3.674176776429809, + "submission": 3.583036025060912, + "synchronized_e2e": 2.655010801771347 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.843016, + "after_init_synchronized_e2e_ms_per_call": 7.874504, + "including_init_host_enqueue_ms_per_call": 42.29246, + "including_init_synchronized_e2e_ms_per_call": 491.915998, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9325361000000001, + "after_init_synchronized_e2e_ms_per_call": 0.9539144, + "including_init_host_enqueue_ms_per_call": 4.3774805, + "including_init_synchronized_e2e_ms_per_call": 49.358063800000004, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.24148811000000003, + "after_init_synchronized_e2e_ms_per_call": 0.26185544, + "including_init_host_enqueue_ms_per_call": 0.58598255, + "including_init_synchronized_e2e_ms_per_call": 5.10227038, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.17238331099999998, + "after_init_synchronized_e2e_ms_per_call": 0.19264954400000003, + "including_init_host_enqueue_ms_per_call": 0.20683275499999998, + "including_init_synchronized_e2e_ms_per_call": 0.6766910380000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.843016, + "synchronized_e2e_ms": 7.874504, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 1.055137, + "median": 0.06784, + "min": 0.057856, + "p90": 0.0892928, + "sample_count": 1779 + }, + "host_enqueue_ms": { + "max": 2.376418, + "median": 0.164705, + "min": 0.134304, + "p90": 0.218829, + "sample_count": 1779 + }, + "sample_count": 1779, + "synchronized_e2e_ms": { + "max": 2.590339, + "median": 0.18496, + "min": 0.1504, + "p90": 0.24330240000000003, + "sample_count": 1779 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 5546, + "candidate_precomputed_gpu_span_ms": 0.018112, + "candidate_precomputed_host_enqueue_ms": 0.042304, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.018112, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.042304, + "candidate_precomputed_synchronized_e2e_ms": 0.054272, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.018112 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.047616, + "synchronized_e2e_ms": 0.064096 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.018112 + }, + "host_enqueue_ms": { + "median": 0.042304 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.018112 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5546, + "submission_ms": { + "median": 0.042304 + }, + "synchronized_e2e_ms": { + "median": 0.054272 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc0073f260", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc0073f320" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 3.1219081272084805, + "submission": 1.2465960665658093, + "synchronized_e2e": 2.0058962264150946 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 1779, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.056544, + "candidate_public_raw_host_enqueue_ms": 0.052736, + "candidate_public_raw_inter_kernel_gap_ms": 0.00016, + "candidate_public_raw_kernel_sum_ms": 0.056384, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.052736, + "candidate_public_raw_synchronized_e2e_ms": 0.108864, + "candidate_public_raw_tflops_from_gpu_span": 77.1448104131296, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.056384 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.260864, + "synchronized_e2e_ms": 0.303712 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.056544 + }, + "host_enqueue_ms": { + "median": 0.052736 + }, + "inter_kernel_gap_ms": { + "median": 0.00016 + }, + "kernel_sum_ms": { + "median": 0.056384 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1779, + "submission_ms": { + "median": 0.052736 + }, + "synchronized_e2e_ms": { + "median": 0.108864 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.501347, + "after_init_synchronized_e2e_ms_per_call": 2.535171, + "including_init_host_enqueue_ms_per_call": 37.375431, + "including_init_synchronized_e2e_ms_per_call": 37.478246999999996, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2975971, + "after_init_synchronized_e2e_ms_per_call": 0.35149470000000005, + "including_init_host_enqueue_ms_per_call": 3.7850054999999996, + "including_init_synchronized_e2e_ms_per_call": 3.8458023, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07722211, + "after_init_synchronized_e2e_ms_per_call": 0.13312707, + "including_init_host_enqueue_ms_per_call": 0.42596294999999995, + "including_init_synchronized_e2e_ms_per_call": 0.48255782999999997, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.055184611, + "after_init_synchronized_e2e_ms_per_call": 0.11129030700000002, + "including_init_host_enqueue_ms_per_call": 0.090058695, + "including_init_synchronized_e2e_ms_per_call": 0.146233383, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.501347, + "synchronized_e2e_ms": 2.535171, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.057376, + "median": 0.056544, + "min": 0.055872, + "p90": 0.056865, + "sample_count": 1779 + }, + "host_enqueue_ms": { + "max": 1.63997, + "median": 0.052736, + "min": 0.039488, + "p90": 0.07453440000000001, + "sample_count": 1779 + }, + "sample_count": 1779, + "synchronized_e2e_ms": { + "max": 1.773026, + "median": 0.108864, + "min": 0.095456, + "p90": 0.1273026, + "sample_count": 1779 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.083296, + "submission_ms": 0.083296, + "synchronized_e2e_ms": 0.103328 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.065696, + "submission_ms": 0.065696, + "synchronized_e2e_ms": 0.08768 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.24592, + "submission_ms": 0.24592, + "synchronized_e2e_ms": 0.269152 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.843016, + "submission_ms": 7.843016, + "synchronized_e2e_ms": 7.874504 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.070112, + "submission_ms": 0.070112, + "synchronized_e2e_ms": 0.089408 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.047616, + "submission_ms": 0.047616, + "synchronized_e2e_ms": 0.064096 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.285506, + "submission_ms": 1.285506, + "synchronized_e2e_ms": 1.310594 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.120289, + "submission_ms": 1.120289, + "synchronized_e2e_ms": 1.144609 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.260864, + "submission_ms": 0.260864, + "synchronized_e2e_ms": 0.303712 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.501347, + "submission_ms": 2.501347, + "synchronized_e2e_ms": 2.535171 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.297119, + "evolution_kernel_ms": 0.155872, + "evolution_speedup": 1.9062, + "evolution_tflops": 27.985, + "expected_route": "aligned_weave_v10_fallback", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_5600_d128_fallback_neighbor_b8_n8320_k256_d128", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 14650, + "measurement_schedule_sha256": "1a145266af9fd920779082a6a3dbbcf93183e46d68bf8476548aa218e10f1ab0", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 5546, + "public_pair_count": 1779, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 5546, + "baseline_public_raw": 1779, + "candidate_precomputed": 5546, + "candidate_public_raw": 1779 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2932 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:adjacent_5600_d128_fallback_neighbor_b8_n8320_k256_d128", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.0194346289752652, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.699000587889477, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.106103690835845, + "including_init_synchronized_e2e_speedup": 13.125373713450367, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.7138798963398307, + "including_init_synchronized_e2e_speedup": 12.834269665916006, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 1.966958635835672, + "including_init_synchronized_e2e_speedup": 10.573386364904701, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.7310541159707646, + "including_init_synchronized_e2e_speedup": 4.627473044236418, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.1997736276174307, + "hot_synchronized_e2e_speedup": 1.699000587889477, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 5601281, + "selected_route": "aligned_weave_v10_fallback", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_5600_d128_fallback_neighbor_b8_n8320_k256_d128", + "source": "guard_miss_fallback", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 3, + "D": 224, + "K": 512, + "N": 3072, + "baseline_07cf_adapter_bench_iters": 3710, + "baseline_07cf_adapter_gpu_span_ms": 0.058464, + "baseline_07cf_adapter_host_enqueue_ms": 0.147088, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.029408, + "baseline_07cf_adapter_kernel_sum_ms": 0.028992, + "baseline_07cf_adapter_submission_ms": 0.147088, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.16532799999999997, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.028992 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.176192, + "synchronized_e2e_ms": 0.194176 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.058464 + }, + "host_enqueue_ms": { + "median": 0.147088 + }, + "inter_kernel_gap_ms": { + "median": 0.029408 + }, + "kernel_sum_ms": { + "median": 0.028992 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3710, + "submission_ms": { + "median": 0.147088 + }, + "synchronized_e2e_ms": { + "median": 0.16532799999999997 + } + }, + "baseline_07cf_precomputed_bench_iters": 5517, + "baseline_07cf_precomputed_gpu_span_ms": 0.018048, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042368, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.018048, + "baseline_07cf_precomputed_submission_ms": 0.042368, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.065888, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.018048 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.037728, + "synchronized_e2e_ms": 0.055648 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.018048 + }, + "host_enqueue_ms": { + "median": 0.042368 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.018048 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5517, + "submission_ms": { + "median": 0.042368 + }, + "synchronized_e2e_ms": { + "median": 0.065888 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 3.2393617021276593, + "submission": 3.4716767371601205, + "synchronized_e2e": 2.5092277804759586 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.091751, + "after_init_synchronized_e2e_ms_per_call": 7.118087, + "including_init_host_enqueue_ms_per_call": 42.643019, + "including_init_synchronized_e2e_ms_per_call": 42.749324, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8415543, + "after_init_synchronized_e2e_ms_per_call": 0.8606038999999999, + "including_init_host_enqueue_ms_per_call": 4.3966811, + "including_init_synchronized_e2e_ms_per_call": 4.4237276, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.21653463000000003, + "after_init_synchronized_e2e_ms_per_call": 0.23485558999999995, + "including_init_host_enqueue_ms_per_call": 0.57204731, + "including_init_synchronized_e2e_ms_per_call": 0.5911679599999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.154032663, + "after_init_synchronized_e2e_ms_per_call": 0.17228075899999998, + "including_init_host_enqueue_ms_per_call": 0.189583931, + "including_init_synchronized_e2e_ms_per_call": 0.207911996, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.091751, + "synchronized_e2e_ms": 7.118087, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.240934, + "median": 0.058464, + "min": 0.049696, + "p90": 0.0813824, + "sample_count": 3710 + }, + "host_enqueue_ms": { + "max": 36.010341, + "median": 0.147088, + "min": 0.119968, + "p90": 0.22586889999999998, + "sample_count": 3710 + }, + "sample_count": 3710, + "synchronized_e2e_ms": { + "max": 36.047397, + "median": 0.16532799999999997, + "min": 0.134816, + "p90": 0.25015049999999994, + "sample_count": 3710 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 5517, + "candidate_precomputed_gpu_span_ms": 0.015776, + "candidate_precomputed_host_enqueue_ms": 0.053376, + "candidate_precomputed_inter_kernel_gap_ms": 0.003168, + "candidate_precomputed_kernel_sum_ms": 0.012576, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.053376, + "candidate_precomputed_synchronized_e2e_ms": 0.0648, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.012576 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.035776, + "synchronized_e2e_ms": 0.049024 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.015776 + }, + "host_enqueue_ms": { + "median": 0.053376 + }, + "inter_kernel_gap_ms": { + "median": 0.003168 + }, + "kernel_sum_ms": { + "median": 0.012576 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5517, + "submission_ms": { + "median": 0.053376 + }, + "synchronized_e2e_ms": { + "median": 0.0648 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01233710", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01230530" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.7119675456389454, + "submission": 0.8330335731414868, + "synchronized_e2e": 1.1125925925925926 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 3710, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.027008, + "candidate_public_raw_host_enqueue_ms": 0.044464000000000004, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.026976, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.044464000000000004, + "candidate_public_raw_synchronized_e2e_ms": 0.072096, + "candidate_public_raw_tflops_from_gpu_span": 78.27048341232226, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.026976 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.06496, + "synchronized_e2e_ms": 0.081696 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.027008 + }, + "host_enqueue_ms": { + "median": 0.044464000000000004 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.026976 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3710, + "submission_ms": { + "median": 0.044464000000000004 + }, + "synchronized_e2e_ms": { + "median": 0.072096 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.264194, + "after_init_synchronized_e2e_ms_per_call": 2.289922, + "including_init_host_enqueue_ms_per_call": 38.102214000000004, + "including_init_synchronized_e2e_ms_per_call": 486.390135, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.266437, + "after_init_synchronized_e2e_ms_per_call": 0.29387859999999993, + "including_init_host_enqueue_ms_per_call": 3.8502390000000006, + "including_init_synchronized_e2e_ms_per_call": 48.703899899999996, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06666129999999999, + "after_init_synchronized_e2e_ms_per_call": 0.09427425999999998, + "including_init_host_enqueue_ms_per_call": 0.4250415, + "including_init_synchronized_e2e_ms_per_call": 4.935276389999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.04668373, + "after_init_synchronized_e2e_ms_per_call": 0.07431382599999999, + "including_init_host_enqueue_ms_per_call": 0.08252174999999999, + "including_init_synchronized_e2e_ms_per_call": 0.558414039, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.264194, + "synchronized_e2e_ms": 2.289922, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.027648, + "median": 0.027008, + "min": 0.026624, + "p90": 0.027232, + "sample_count": 3710 + }, + "host_enqueue_ms": { + "max": 0.501985, + "median": 0.044464000000000004, + "min": 0.035904, + "p90": 0.07127679999999997, + "sample_count": 3710 + }, + "sample_count": 3710, + "synchronized_e2e_ms": { + "max": 0.559073, + "median": 0.072096, + "min": 0.065024, + "p90": 0.09499519999999999, + "sample_count": 3710 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.060224, + "submission_ms": 0.060224, + "synchronized_e2e_ms": 0.07664 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.037728, + "submission_ms": 0.037728, + "synchronized_e2e_ms": 0.055648 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.176192, + "submission_ms": 0.176192, + "synchronized_e2e_ms": 0.194176 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.091751, + "submission_ms": 7.091751, + "synchronized_e2e_ms": 7.118087 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.064512, + "submission_ms": 0.064512, + "synchronized_e2e_ms": 0.078752 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.035776, + "submission_ms": 0.035776, + "synchronized_e2e_ms": 0.049024 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.191809, + "submission_ms": 1.191809, + "synchronized_e2e_ms": 1.231585 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.032354, + "submission_ms": 1.032354, + "synchronized_e2e_ms": 1.056482 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.06496, + "submission_ms": 0.06496, + "synchronized_e2e_ms": 0.081696 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.264194, + "submission_ms": 2.264194, + "synchronized_e2e_ms": 2.289922 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.343839, + "evolution_kernel_ms": 0.174928, + "evolution_speedup": 1.9656, + "evolution_tflops": 12.0846, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_5600_d224_k512_overlap_b3_n3072_k512_d224", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 18454, + "measurement_schedule_sha256": "ecc3b75d6545b83513651a0e4742aff919007b355f1117692e9a28afaa35f57b", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 5517, + "public_pair_count": 3710, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 5517, + "baseline_public_raw": 3710, + "candidate_precomputed": 5517, + "candidate_public_raw": 3710 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 3692 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:adjacent_5600_d224_k512_overlap_b3_n3072_k512_d224", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.1440162271805276, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.2931646693297822, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.1084408115210915, + "including_init_synchronized_e2e_speedup": 0.08789101777321204, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.928433373508653, + "including_init_synchronized_e2e_speedup": 0.09082902209233558, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.4911952636912766, + "including_init_synchronized_e2e_speedup": 0.11978416471220166, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.3182867613356364, + "including_init_synchronized_e2e_speedup": 0.3723258755677523, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.1646919431279623, + "hot_synchronized_e2e_speedup": 2.2931646693297822, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 5602241, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_5600_d224_k512_overlap_b3_n3072_k512_d224", + "source": "guard_overlap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 288, + "K": 2048, + "N": 1024, + "baseline_07cf_adapter_bench_iters": 1422, + "baseline_07cf_adapter_gpu_span_ms": 0.101008, + "baseline_07cf_adapter_host_enqueue_ms": 0.15732800000000002, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.037823999999999997, + "baseline_07cf_adapter_kernel_sum_ms": 0.063136, + "baseline_07cf_adapter_submission_ms": 0.15732800000000002, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.207008, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.063136 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.21376, + "synchronized_e2e_ms": 0.262016 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.101008 + }, + "host_enqueue_ms": { + "median": 0.15732800000000002 + }, + "inter_kernel_gap_ms": { + "median": 0.037823999999999997 + }, + "kernel_sum_ms": { + "median": 0.063136 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1422, + "submission_ms": { + "median": 0.15732800000000002 + }, + "synchronized_e2e_ms": { + "median": 0.207008 + } + }, + "baseline_07cf_precomputed_bench_iters": 1589, + "baseline_07cf_precomputed_gpu_span_ms": 0.063104, + "baseline_07cf_precomputed_host_enqueue_ms": 0.046816, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.063104, + "baseline_07cf_precomputed_submission_ms": 0.046816, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.116736, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.063104 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.048928, + "synchronized_e2e_ms": 0.103712 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.063104 + }, + "host_enqueue_ms": { + "median": 0.046816 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.063104 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1589, + "submission_ms": { + "median": 0.046816 + }, + "synchronized_e2e_ms": { + "median": 0.116736 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.6006592292089252, + "submission": 3.36056049213944, + "synchronized_e2e": 1.773300438596491 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.808456, + "after_init_synchronized_e2e_ms_per_call": 7.837256, + "including_init_host_enqueue_ms_per_call": 44.742894, + "including_init_synchronized_e2e_ms_per_call": 462.94857500000006, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9224408000000001, + "after_init_synchronized_e2e_ms_per_call": 0.9700328, + "including_init_host_enqueue_ms_per_call": 4.615884599999999, + "including_init_synchronized_e2e_ms_per_call": 46.48116470000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23383928000000004, + "after_init_synchronized_e2e_ms_per_call": 0.28331048, + "including_init_host_enqueue_ms_per_call": 0.60318366, + "including_init_synchronized_e2e_ms_per_call": 4.8344236700000005, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.16497912800000003, + "after_init_synchronized_e2e_ms_per_call": 0.214638248, + "including_init_host_enqueue_ms_per_call": 0.20191356600000004, + "including_init_synchronized_e2e_ms_per_call": 0.6697495670000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.808456, + "synchronized_e2e_ms": 7.837256, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.406401, + "median": 0.101008, + "min": 0.090464, + "p90": 0.1102048, + "sample_count": 1422 + }, + "host_enqueue_ms": { + "max": 0.667553, + "median": 0.15732800000000002, + "min": 0.13104, + "p90": 0.1804289, + "sample_count": 1422 + }, + "sample_count": 1422, + "synchronized_e2e_ms": { + "max": 0.712129, + "median": 0.207008, + "min": 0.184065, + "p90": 0.2286912, + "sample_count": 1422 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 1589, + "candidate_precomputed_gpu_span_ms": 0.076192, + "candidate_precomputed_host_enqueue_ms": 0.042688, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.076192, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.042688, + "candidate_precomputed_synchronized_e2e_ms": 0.113088, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.076192 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.038336, + "synchronized_e2e_ms": 0.083968 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.076192 + }, + "host_enqueue_ms": { + "median": 0.042688 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.076192 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1589, + "submission_ms": { + "median": 0.042688 + }, + "synchronized_e2e_ms": { + "median": 0.113088 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc03ca0bc0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc03ca0d70" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 0.9252414951700967, + "submission": 1.1900299850074962, + "synchronized_e2e": 1.0841822297679684 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 1422, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.070496, + "candidate_public_raw_host_enqueue_ms": 0.0508, + "candidate_public_raw_inter_kernel_gap_ms": 0.000192, + "candidate_public_raw_kernel_sum_ms": 0.070304, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.0508, + "candidate_public_raw_synchronized_e2e_ms": 0.122608, + "candidate_public_raw_tflops_from_gpu_span": 34.270300499319106, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.070304 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.307041, + "synchronized_e2e_ms": 0.361953 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.070496 + }, + "host_enqueue_ms": { + "median": 0.0508 + }, + "inter_kernel_gap_ms": { + "median": 0.000192 + }, + "kernel_sum_ms": { + "median": 0.070304 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1422, + "submission_ms": { + "median": 0.0508 + }, + "synchronized_e2e_ms": { + "median": 0.122608 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.305347, + "after_init_synchronized_e2e_ms_per_call": 3.335235, + "including_init_host_enqueue_ms_per_call": 40.66336999999999, + "including_init_synchronized_e2e_ms_per_call": 40.770858, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.37625469999999994, + "after_init_synchronized_e2e_ms_per_call": 0.4438707, + "including_init_host_enqueue_ms_per_call": 4.112056999999999, + "including_init_synchronized_e2e_ms_per_call": 4.187433, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.08334546999999999, + "after_init_synchronized_e2e_ms_per_call": 0.15473427, + "including_init_host_enqueue_ms_per_call": 0.4569256999999999, + "including_init_synchronized_e2e_ms_per_call": 0.5290904999999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.054054546999999994, + "after_init_synchronized_e2e_ms_per_call": 0.125820627, + "including_init_host_enqueue_ms_per_call": 0.09141256999999998, + "including_init_synchronized_e2e_ms_per_call": 0.16325625, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.305347, + "synchronized_e2e_ms": 3.335235, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.071616, + "median": 0.070496, + "min": 0.069984, + "p90": 0.070945, + "sample_count": 1422 + }, + "host_enqueue_ms": { + "max": 0.359904, + "median": 0.0508, + "min": 0.03888, + "p90": 0.06395840000000001, + "sample_count": 1422 + }, + "sample_count": 1422, + "synchronized_e2e_ms": { + "max": 0.427424, + "median": 0.122608, + "min": 0.113088, + "p90": 0.1345289, + "sample_count": 1422 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.10704, + "submission_ms": 0.10704, + "synchronized_e2e_ms": 0.15952 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.048928, + "submission_ms": 0.048928, + "synchronized_e2e_ms": 0.103712 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.21376, + "submission_ms": 0.21376, + "synchronized_e2e_ms": 0.262016 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.808456, + "submission_ms": 7.808456, + "synchronized_e2e_ms": 7.837256 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.093344, + "submission_ms": 0.093344, + "synchronized_e2e_ms": 0.111392 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.038336, + "submission_ms": 0.038336, + "synchronized_e2e_ms": 0.083968 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.487361, + "submission_ms": 1.487361, + "synchronized_e2e_ms": 1.507585 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.205154, + "submission_ms": 1.205154, + "synchronized_e2e_ms": 1.226562 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.307041, + "submission_ms": 0.307041, + "synchronized_e2e_ms": 0.361953 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.305347, + "submission_ms": 3.305347, + "synchronized_e2e_ms": 3.335235 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.976096, + "evolution_kernel_ms": 0.206512, + "evolution_speedup": 4.7266, + "evolution_tflops": 11.6987, + "expected_route": "d288_parent_splitk_hybrid_20260629_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_5600_d288_mid_highk_b2_n1024_k2048_d288", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 6022, + "measurement_schedule_sha256": "4aa62eb27325bfa3345d02fbaa8edc3ac21b895bb8e2ef5238f42e63b14ab45b", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 1589, + "public_pair_count": 1422, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 1589, + "baseline_public_raw": 1422, + "candidate_precomputed": 1589, + "candidate_public_raw": 1422 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 1206 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:adjacent_5600_d288_mid_highk_b2_n1024_k2048_d288", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.8282234355312893, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.6883726999869504, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.3498362184373813, + "including_init_synchronized_e2e_speedup": 11.354889195611241, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.185394981015868, + "including_init_synchronized_e2e_speedup": 11.100157232366465, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 1.8309485028752839, + "including_init_synchronized_e2e_speedup": 9.137233932569194, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.7059066793555242, + "including_init_synchronized_e2e_speedup": 4.102443655296505, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.4328188833408988, + "hot_synchronized_e2e_speedup": 1.6883726999869504, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 5602881, + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_5600_d288_mid_highk_b2_n1024_k2048_d288", + "source": "heldout_neighborhood", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 352, + "K": 4096, + "N": 768, + "baseline_07cf_adapter_bench_iters": 1901, + "baseline_07cf_adapter_gpu_span_ms": 0.148993, + "baseline_07cf_adapter_host_enqueue_ms": 0.147872, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.027296, + "baseline_07cf_adapter_kernel_sum_ms": 0.121664, + "baseline_07cf_adapter_submission_ms": 0.147872, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.250144, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.121664 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.179872, + "synchronized_e2e_ms": 0.281152 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.148993 + }, + "host_enqueue_ms": { + "median": 0.147872 + }, + "inter_kernel_gap_ms": { + "median": 0.027296 + }, + "kernel_sum_ms": { + "median": 0.121664 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1901, + "submission_ms": { + "median": 0.147872 + }, + "synchronized_e2e_ms": { + "median": 0.250144 + } + }, + "baseline_07cf_precomputed_bench_iters": 2446, + "baseline_07cf_precomputed_gpu_span_ms": 0.125312, + "baseline_07cf_precomputed_host_enqueue_ms": 0.043056, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.125312, + "baseline_07cf_precomputed_submission_ms": 0.043056, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.174816, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.125312 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.062592, + "synchronized_e2e_ms": 0.161632 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.125312 + }, + "host_enqueue_ms": { + "median": 0.043056 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.125312 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2446, + "submission_ms": { + "median": 0.043056 + }, + "synchronized_e2e_ms": { + "median": 0.174816 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.1889763151174666, + "submission": 3.4344109996283914, + "synchronized_e2e": 1.4308987735676366 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.90532, + "after_init_synchronized_e2e_ms_per_call": 7.988744, + "including_init_host_enqueue_ms_per_call": 42.133227000000005, + "including_init_synchronized_e2e_ms_per_call": 42.334187, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9236167999999999, + "after_init_synchronized_e2e_ms_per_call": 1.0240040000000001, + "including_init_host_enqueue_ms_per_call": 4.346407500000001, + "including_init_synchronized_e2e_ms_per_call": 4.4585482999999995, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22544648000000003, + "after_init_synchronized_e2e_ms_per_call": 0.32752999999999993, + "including_init_host_enqueue_ms_per_call": 0.56772555, + "including_init_synchronized_e2e_ms_per_call": 0.6709844300000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.155629448, + "after_init_synchronized_e2e_ms_per_call": 0.25788259999999996, + "including_init_host_enqueue_ms_per_call": 0.189857355, + "including_init_synchronized_e2e_ms_per_call": 0.29222804299999994, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.90532, + "synchronized_e2e_ms": 7.988744, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 8.831305, + "median": 0.148993, + "min": 0.140992, + "p90": 0.16912, + "sample_count": 1901 + }, + "host_enqueue_ms": { + "max": 35.600325, + "median": 0.147872, + "min": 0.124, + "p90": 0.221792, + "sample_count": 1901 + }, + "sample_count": 1901, + "synchronized_e2e_ms": { + "max": 35.670725, + "median": 0.250144, + "min": 0.22912, + "p90": 0.315392, + "sample_count": 1901 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 2446, + "candidate_precomputed_gpu_span_ms": 0.041088, + "candidate_precomputed_host_enqueue_ms": 0.049952, + "candidate_precomputed_inter_kernel_gap_ms": 0.00224, + "candidate_precomputed_kernel_sum_ms": 0.038784, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.049952, + "candidate_precomputed_synchronized_e2e_ms": 0.07576, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.038784 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.036608, + "synchronized_e2e_ms": 0.057056 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.041088 + }, + "host_enqueue_ms": { + "median": 0.049952 + }, + "inter_kernel_gap_ms": { + "median": 0.00224 + }, + "kernel_sum_ms": { + "median": 0.038784 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2446, + "submission_ms": { + "median": 0.049952 + }, + "synchronized_e2e_ms": { + "median": 0.07576 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04536f30", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295367b0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.280373831775701, + "submission": 0.9423446508648302, + "synchronized_e2e": 1.340232312565998 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 1901, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.052608, + "candidate_public_raw_host_enqueue_ms": 0.047072, + "candidate_public_raw_inter_kernel_gap_ms": 0.000128, + "candidate_public_raw_kernel_sum_ms": 0.05248, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.047072, + "candidate_public_raw_synchronized_e2e_ms": 0.101536, + "candidate_public_raw_tflops_from_gpu_span": 84.19223357664234, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.05248 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.254752, + "synchronized_e2e_ms": 0.290625 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.052608 + }, + "host_enqueue_ms": { + "median": 0.047072 + }, + "inter_kernel_gap_ms": { + "median": 0.000128 + }, + "kernel_sum_ms": { + "median": 0.05248 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1901, + "submission_ms": { + "median": 0.047072 + }, + "synchronized_e2e_ms": { + "median": 0.101536 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.744067, + "after_init_synchronized_e2e_ms_per_call": 2.768163, + "including_init_host_enqueue_ms_per_call": 37.321607, + "including_init_synchronized_e2e_ms_per_call": 450.855955, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.3167715, + "after_init_synchronized_e2e_ms_per_call": 0.3681987, + "including_init_host_enqueue_ms_per_call": 3.7745255, + "including_init_synchronized_e2e_ms_per_call": 45.1769779, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07404195, + "after_init_synchronized_e2e_ms_per_call": 0.12820226999999998, + "including_init_host_enqueue_ms_per_call": 0.41981735, + "including_init_synchronized_e2e_ms_per_call": 4.609080189999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.049768995, + "after_init_synchronized_e2e_ms_per_call": 0.104202627, + "including_init_host_enqueue_ms_per_call": 0.084346535, + "including_init_synchronized_e2e_ms_per_call": 0.5522904190000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.744067, + "synchronized_e2e_ms": 2.768163, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.054272, + "median": 0.052608, + "min": 0.051456, + "p90": 0.053248, + "sample_count": 1901 + }, + "host_enqueue_ms": { + "max": 50.918709, + "median": 0.047072, + "min": 0.037408, + "p90": 0.07136, + "sample_count": 1901 + }, + "sample_count": 1901, + "synchronized_e2e_ms": { + "max": 51.179253, + "median": 0.101536, + "min": 0.093568, + "p90": 0.121856, + "sample_count": 1901 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.064384, + "submission_ms": 0.064384, + "synchronized_e2e_ms": 0.162848 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.062592, + "submission_ms": 0.062592, + "synchronized_e2e_ms": 0.161632 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.179872, + "submission_ms": 0.179872, + "synchronized_e2e_ms": 0.281152 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.90532, + "submission_ms": 7.90532, + "synchronized_e2e_ms": 7.988744 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.06064, + "submission_ms": 0.06064, + "synchronized_e2e_ms": 0.075296 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.036608, + "submission_ms": 0.036608, + "synchronized_e2e_ms": 0.057056 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.316385, + "submission_ms": 1.316385, + "synchronized_e2e_ms": 1.336577 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.271394, + "submission_ms": 1.271394, + "synchronized_e2e_ms": 1.29165 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.254752, + "submission_ms": 0.254752, + "synchronized_e2e_ms": 0.290625 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.744067, + "submission_ms": 2.744067, + "synchronized_e2e_ms": 2.768163 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 1.621022, + "evolution_kernel_ms": 1.051615, + "evolution_speedup": 1.5415, + "evolution_tflops": 4.2118, + "expected_route": "d352_exactd_splitk_c95c_v2", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_5600_d352_splitk_edge_b2_n768_k4096_d352", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 8694, + "measurement_schedule_sha256": "85192f24ee6b17b6f56fc1978fd49d9c8dd1b6812df0bd7476efed562d155b2e", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 2446, + "public_pair_count": 1901, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 2446, + "baseline_public_raw": 1901, + "candidate_precomputed": 2446, + "candidate_public_raw": 1901 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 1742 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:adjacent_5600_d352_splitk_edge_b2_n768_k4096_d352", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 3.0498442367601246, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.463599117554365, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.8859369914271666, + "including_init_synchronized_e2e_speedup": 0.09389736684303084, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.7811179126922507, + "including_init_synchronized_e2e_speedup": 0.09869071609590778, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.554790956509584, + "including_init_synchronized_e2e_speedup": 0.14557881450094715, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.4748186051010013, + "including_init_synchronized_e2e_speedup": 0.5291202471502586, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.8321357968369827, + "hot_synchronized_e2e_speedup": 2.463599117554365, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 5603521, + "selected_route": "d352_exactd_splitk_c95c_v2", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_5600_d352_splitk_edge_b2_n768_k4096_d352", + "source": "request_specific", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 416, + "K": 256, + "N": 2048, + "baseline_07cf_adapter_bench_iters": 3180, + "baseline_07cf_adapter_gpu_span_ms": 0.058624, + "baseline_07cf_adapter_host_enqueue_ms": 0.155312, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.031392, + "baseline_07cf_adapter_kernel_sum_ms": 0.0272, + "baseline_07cf_adapter_submission_ms": 0.155312, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.175072, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.0272 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.247744, + "synchronized_e2e_ms": 0.273345 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.058624 + }, + "host_enqueue_ms": { + "median": 0.155312 + }, + "inter_kernel_gap_ms": { + "median": 0.031392 + }, + "kernel_sum_ms": { + "median": 0.0272 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3180, + "submission_ms": { + "median": 0.155312 + }, + "synchronized_e2e_ms": { + "median": 0.175072 + } + }, + "baseline_07cf_precomputed_bench_iters": 7015, + "baseline_07cf_precomputed_gpu_span_ms": 0.014336, + "baseline_07cf_precomputed_host_enqueue_ms": 0.044064, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.014336, + "baseline_07cf_precomputed_submission_ms": 0.044064, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.063968, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.014336 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.055616, + "synchronized_e2e_ms": 0.07328 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.014336 + }, + "host_enqueue_ms": { + "median": 0.044064 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.014336 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 7015, + "submission_ms": { + "median": 0.044064 + }, + "synchronized_e2e_ms": { + "median": 0.063968 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 4.089285714285714, + "submission": 3.5246913580246915, + "synchronized_e2e": 2.736868434217109 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.401576, + "after_init_synchronized_e2e_ms_per_call": 7.434184, + "including_init_host_enqueue_ms_per_call": 41.85102, + "including_init_synchronized_e2e_ms_per_call": 491.475678, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8799384, + "after_init_synchronized_e2e_ms_per_call": 0.9009832, + "including_init_host_enqueue_ms_per_call": 4.324882799999999, + "including_init_synchronized_e2e_ms_per_call": 49.3051326, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22777464000000003, + "after_init_synchronized_e2e_ms_per_call": 0.24766312, + "including_init_host_enqueue_ms_per_call": 0.57226908, + "including_init_synchronized_e2e_ms_per_call": 5.08807806, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.162558264, + "after_init_synchronized_e2e_ms_per_call": 0.182331112, + "including_init_host_enqueue_ms_per_call": 0.197007708, + "including_init_synchronized_e2e_ms_per_call": 0.666372606, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.401576, + "synchronized_e2e_ms": 7.434184, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.572897, + "median": 0.058624, + "min": 0.048384, + "p90": 0.0816, + "sample_count": 3180 + }, + "host_enqueue_ms": { + "max": 25.270298, + "median": 0.155312, + "min": 0.130048, + "p90": 0.2006432, + "sample_count": 3180 + }, + "sample_count": 3180, + "synchronized_e2e_ms": { + "max": 25.299226, + "median": 0.175072, + "min": 0.145184, + "p90": 0.22514969999999995, + "sample_count": 3180 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 7015, + "candidate_precomputed_gpu_span_ms": 0.01888, + "candidate_precomputed_host_enqueue_ms": 0.054528, + "candidate_precomputed_inter_kernel_gap_ms": 0.002592, + "candidate_precomputed_kernel_sum_ms": 0.016352, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.054528, + "candidate_precomputed_synchronized_e2e_ms": 0.066624, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.016352 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.051264, + "synchronized_e2e_ms": 0.069952 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.01888 + }, + "host_enqueue_ms": { + "median": 0.054528 + }, + "inter_kernel_gap_ms": { + "median": 0.002592 + }, + "kernel_sum_ms": { + "median": 0.016352 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 7015, + "submission_ms": { + "median": 0.054528 + }, + "synchronized_e2e_ms": { + "median": 0.066624 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295bf110", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04dca1e0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.694915254237288, + "submission": 0.8870305164319249, + "synchronized_e2e": 1.2041306436119117 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 3180, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.032, + "candidate_public_raw_host_enqueue_ms": 0.048368, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.031936, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.048368, + "candidate_public_raw_synchronized_e2e_ms": 0.080224, + "candidate_public_raw_tflops_from_gpu_span": 54.525952, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.031936 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.062944, + "synchronized_e2e_ms": 0.084448 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.032 + }, + "host_enqueue_ms": { + "median": 0.048368 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.031936 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3180, + "submission_ms": { + "median": 0.048368 + }, + "synchronized_e2e_ms": { + "median": 0.080224 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.490402, + "after_init_synchronized_e2e_ms_per_call": 2.522658, + "including_init_host_enqueue_ms_per_call": 37.364486, + "including_init_synchronized_e2e_ms_per_call": 37.465734, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.29257140000000004, + "after_init_synchronized_e2e_ms_per_call": 0.32446739999999996, + "including_init_host_enqueue_ms_per_call": 3.7799798000000004, + "including_init_synchronized_e2e_ms_per_call": 3.8187749999999996, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07278833999999999, + "after_init_synchronized_e2e_ms_per_call": 0.10464833999999999, + "including_init_host_enqueue_ms_per_call": 0.42152918, + "including_init_synchronized_e2e_ms_per_call": 0.4540791, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.050810034000000004, + "after_init_synchronized_e2e_ms_per_call": 0.08266643400000001, + "including_init_host_enqueue_ms_per_call": 0.085684118, + "including_init_synchronized_e2e_ms_per_call": 0.11760951, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.490402, + "synchronized_e2e_ms": 2.522658, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.032704, + "median": 0.032, + "min": 0.031168, + "p90": 0.032256, + "sample_count": 3180 + }, + "host_enqueue_ms": { + "max": 23.139128, + "median": 0.048368, + "min": 0.03968, + "p90": 0.070208, + "sample_count": 3180 + }, + "sample_count": 3180, + "synchronized_e2e_ms": { + "max": 28.253149, + "median": 0.080224, + "min": 0.072256, + "p90": 0.10007049999999999, + "sample_count": 3180 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.07856, + "submission_ms": 0.07856, + "synchronized_e2e_ms": 0.099936 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.055616, + "submission_ms": 0.055616, + "synchronized_e2e_ms": 0.07328 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.247744, + "submission_ms": 0.247744, + "synchronized_e2e_ms": 0.273345 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.401576, + "submission_ms": 7.401576, + "synchronized_e2e_ms": 7.434184 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.077216, + "submission_ms": 0.077216, + "synchronized_e2e_ms": 0.097088 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.051264, + "submission_ms": 0.051264, + "synchronized_e2e_ms": 0.069952 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.525825, + "submission_ms": 1.525825, + "synchronized_e2e_ms": 1.557057 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.241506, + "submission_ms": 1.241506, + "synchronized_e2e_ms": 1.26989 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.062944, + "submission_ms": 0.062944, + "synchronized_e2e_ms": 0.084448 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.490402, + "submission_ms": 2.490402, + "synchronized_e2e_ms": 2.522658 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.370128, + "evolution_kernel_ms": 0.175743, + "evolution_speedup": 2.1061, + "evolution_tflops": 9.9283, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_5600_d416_smallk_boundary_b4_n2048_k256_d416", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 20390, + "measurement_schedule_sha256": "dbf2118b3b542c338a56e76354647d1a62faf4a2bbd2d1e560442f62e501182d", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 7015, + "public_pair_count": 3180, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 7015, + "baseline_public_raw": 3180, + "candidate_precomputed": 7015, + "candidate_public_raw": 3180 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 4078 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:adjacent_5600_d416_smallk_boundary_b4_n2048_k256_d416", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.759322033898305, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.182289589150379, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.946964669804627, + "including_init_synchronized_e2e_speedup": 13.118004788055135, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.776806545125951, + "including_init_synchronized_e2e_speedup": 12.911243160437577, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.366622537920812, + "including_init_synchronized_e2e_speedup": 11.205268112978553, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.2056244980883046, + "including_init_synchronized_e2e_speedup": 5.665975532080696, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.832, + "hot_synchronized_e2e_speedup": 2.182289589150379, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 5604161, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_5600_d416_smallk_boundary_b4_n2048_k256_d416", + "source": "guard_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 480, + "K": 1024, + "N": 1536, + "baseline_07cf_adapter_bench_iters": 4088, + "baseline_07cf_adapter_gpu_span_ms": 0.075968, + "baseline_07cf_adapter_host_enqueue_ms": 0.153968, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.038960499999999995, + "baseline_07cf_adapter_kernel_sum_ms": 0.037024, + "baseline_07cf_adapter_submission_ms": 0.153968, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.179888, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.037024 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.283745, + "synchronized_e2e_ms": 0.304609 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.075968 + }, + "host_enqueue_ms": { + "median": 0.153968 + }, + "inter_kernel_gap_ms": { + "median": 0.038960499999999995 + }, + "kernel_sum_ms": { + "median": 0.037024 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 4088, + "submission_ms": { + "median": 0.153968 + }, + "synchronized_e2e_ms": { + "median": 0.179888 + } + }, + "baseline_07cf_precomputed_bench_iters": 4623, + "baseline_07cf_precomputed_gpu_span_ms": 0.03584, + "baseline_07cf_precomputed_host_enqueue_ms": 0.044672, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.03584, + "baseline_07cf_precomputed_submission_ms": 0.044672, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.085728, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.03584 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.044256, + "synchronized_e2e_ms": 0.07744 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.03584 + }, + "host_enqueue_ms": { + "median": 0.044672 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.03584 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4623, + "submission_ms": { + "median": 0.044672 + }, + "synchronized_e2e_ms": { + "median": 0.085728 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.119642857142857, + "submission": 3.446633237822349, + "synchronized_e2e": 2.0983575961179546 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.94292, + "after_init_synchronized_e2e_ms_per_call": 7.970888, + "including_init_host_enqueue_ms_per_call": 43.494188, + "including_init_synchronized_e2e_ms_per_call": 43.602125, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9328631999999999, + "after_init_synchronized_e2e_ms_per_call": 0.9589880000000001, + "including_init_host_enqueue_ms_per_call": 4.48799, + "including_init_synchronized_e2e_ms_per_call": 4.5221117, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23185752, + "after_init_synchronized_e2e_ms_per_call": 0.257798, + "including_init_host_enqueue_ms_per_call": 0.5873702000000001, + "including_init_synchronized_e2e_ms_per_call": 0.61411037, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.16175695199999998, + "after_init_synchronized_e2e_ms_per_call": 0.187679, + "including_init_host_enqueue_ms_per_call": 0.19730822, + "including_init_synchronized_e2e_ms_per_call": 0.223310237, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.94292, + "synchronized_e2e_ms": 7.970888, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.385408, + "median": 0.075968, + "min": 0.065792, + "p90": 0.10229120000000001, + "sample_count": 4088 + }, + "host_enqueue_ms": { + "max": 34.9985, + "median": 0.153968, + "min": 0.123328, + "p90": 0.2533952, + "sample_count": 4088 + }, + "sample_count": 4088, + "synchronized_e2e_ms": { + "max": 35.075524, + "median": 0.179888, + "min": 0.152224, + "p90": 0.2792682, + "sample_count": 4088 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 4623, + "candidate_precomputed_gpu_span_ms": 0.02192, + "candidate_precomputed_host_enqueue_ms": 0.05168, + "candidate_precomputed_inter_kernel_gap_ms": 0.002208, + "candidate_precomputed_kernel_sum_ms": 0.019744, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.05168, + "candidate_precomputed_synchronized_e2e_ms": 0.063904, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.019744 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.03984, + "synchronized_e2e_ms": 0.053056 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.02192 + }, + "host_enqueue_ms": { + "median": 0.05168 + }, + "inter_kernel_gap_ms": { + "median": 0.002208 + }, + "kernel_sum_ms": { + "median": 0.019744 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4623, + "submission_ms": { + "median": 0.05168 + }, + "synchronized_e2e_ms": { + "median": 0.063904 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc029a14c0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc029a2240" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.1197080291970805, + "submission": 0.9430534055727555, + "synchronized_e2e": 1.1529872934401602 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 4088, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.024544, + "candidate_public_raw_host_enqueue_ms": 0.048737, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.02448, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.048737, + "candidate_public_raw_synchronized_e2e_ms": 0.0736805, + "candidate_public_raw_tflops_from_gpu_span": 61.52010430247718, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.02448 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.680481, + "synchronized_e2e_ms": 0.706881 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.024544 + }, + "host_enqueue_ms": { + "median": 0.048737 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.02448 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 4088, + "submission_ms": { + "median": 0.048737 + }, + "synchronized_e2e_ms": { + "median": 0.0736805 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 70.467273, + "after_init_synchronized_e2e_ms_per_call": 70.499753, + "including_init_host_enqueue_ms_per_call": 106.305293, + "including_init_synchronized_e2e_ms_per_call": 554.599966, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 7.0905906000000005, + "after_init_synchronized_e2e_ms_per_call": 7.11628775, + "including_init_host_enqueue_ms_per_call": 10.674392600000001, + "including_init_synchronized_e2e_ms_per_call": 55.526309049999995, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.75292236, + "after_init_synchronized_e2e_ms_per_call": 0.777941225, + "including_init_host_enqueue_ms_per_call": 1.11130256, + "including_init_synchronized_e2e_ms_per_call": 5.618943355, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.119155536, + "after_init_synchronized_e2e_ms_per_call": 0.1441065725, + "including_init_host_enqueue_ms_per_call": 0.154993556, + "including_init_synchronized_e2e_ms_per_call": 0.6282067855, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 70.467273, + "synchronized_e2e_ms": 70.499753, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.025632, + "median": 0.024544, + "min": 0.023776, + "p90": 0.024864, + "sample_count": 4088 + }, + "host_enqueue_ms": { + "max": 36.796166, + "median": 0.048737, + "min": 0.036672, + "p90": 0.0828992, + "sample_count": 4088 + }, + "sample_count": 4088, + "synchronized_e2e_ms": { + "max": 36.886182, + "median": 0.0736805, + "min": 0.06384, + "p90": 0.10582080000000002, + "sample_count": 4088 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.069345, + "submission_ms": 0.069345, + "synchronized_e2e_ms": 0.099873 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.044256, + "submission_ms": 0.044256, + "synchronized_e2e_ms": 0.07744 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.283745, + "submission_ms": 0.283745, + "synchronized_e2e_ms": 0.304609 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.94292, + "submission_ms": 7.94292, + "synchronized_e2e_ms": 7.970888 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.06544, + "submission_ms": 0.06544, + "synchronized_e2e_ms": 0.080096 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.03984, + "submission_ms": 0.03984, + "synchronized_e2e_ms": 0.053056 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.423681, + "submission_ms": 1.423681, + "synchronized_e2e_ms": 1.445217 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.184353, + "submission_ms": 1.184353, + "synchronized_e2e_ms": 1.204449 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.680481, + "submission_ms": 0.680481, + "synchronized_e2e_ms": 0.706881 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 70.467273, + "submission_ms": 70.467273, + "synchronized_e2e_ms": 70.499753 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.718975, + "evolution_kernel_ms": 0.199488, + "evolution_speedup": 3.6041, + "evolution_tflops": 7.5691, + "expected_route": "d480_splitk_k1024_eac2_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_5600_d480_random_b1_n1536_k1024_d480", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 17422, + "measurement_schedule_sha256": "48bfa7612670e892369b752838bb8281d784df1fff906f453cc2f18ac072253b", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 4623, + "public_pair_count": 4088, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 4623, + "baseline_public_raw": 4088, + "candidate_precomputed": 4623, + "candidate_public_raw": 4088 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 3486 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:adjacent_5600_d480_random_b1_n1536_k1024_d480", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.635036496350365, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.441460087811565, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.11306263725491351, + "including_init_synchronized_e2e_speedup": 0.07861905458537298, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.13475958725811785, + "including_init_synchronized_e2e_speedup": 0.08144088410284854, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.33138493206861486, + "including_init_synchronized_e2e_speedup": 0.1092928565392167, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.3023625275661872, + "including_init_synchronized_e2e_speedup": 0.35547250070255726, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 3.0951760104302477, + "hot_synchronized_e2e_speedup": 2.441460087811565, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 5604801, + "selected_route": "d480_splitk_k1024_eac2_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_5600_d480_random_b1_n1536_k1024_d480", + "source": "random_legal", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 3, + "D": 48, + "K": 256, + "N": 1536, + "baseline_07cf_adapter_bench_iters": 8290, + "baseline_07cf_adapter_gpu_span_ms": 0.054496, + "baseline_07cf_adapter_host_enqueue_ms": 0.15968, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.04392, + "baseline_07cf_adapter_kernel_sum_ms": 0.01056, + "baseline_07cf_adapter_submission_ms": 0.15968, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.179904, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.01056 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.236576, + "synchronized_e2e_ms": 0.257312 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.054496 + }, + "host_enqueue_ms": { + "median": 0.15968 + }, + "inter_kernel_gap_ms": { + "median": 0.04392 + }, + "kernel_sum_ms": { + "median": 0.01056 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 8290, + "submission_ms": { + "median": 0.15968 + }, + "synchronized_e2e_ms": { + "median": 0.179904 + } + }, + "baseline_07cf_precomputed_bench_iters": 14502, + "baseline_07cf_precomputed_gpu_span_ms": 0.00688, + "baseline_07cf_precomputed_host_enqueue_ms": 0.045792, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.00688, + "baseline_07cf_precomputed_submission_ms": 0.045792, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.060672, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.00688 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.052736, + "synchronized_e2e_ms": 0.069888 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.00688 + }, + "host_enqueue_ms": { + "median": 0.045792 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.00688 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 14502, + "submission_ms": { + "median": 0.045792 + }, + "synchronized_e2e_ms": { + "median": 0.060672 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 7.92093023255814, + "submission": 3.4870719776380152, + "synchronized_e2e": 2.9651898734177218 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.630088, + "after_init_synchronized_e2e_ms_per_call": 7.659112, + "including_init_host_enqueue_ms_per_call": 44.564526, + "including_init_synchronized_e2e_ms_per_call": 462.77043100000003, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9067207999999999, + "after_init_synchronized_e2e_ms_per_call": 0.9278248000000001, + "including_init_host_enqueue_ms_per_call": 4.6001646, + "including_init_synchronized_e2e_ms_per_call": 46.438956700000006, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23438408, + "after_init_synchronized_e2e_ms_per_call": 0.25469608, + "including_init_host_enqueue_ms_per_call": 0.6037284599999999, + "including_init_synchronized_e2e_ms_per_call": 4.80580927, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.167150408, + "after_init_synchronized_e2e_ms_per_call": 0.187383208, + "including_init_host_enqueue_ms_per_call": 0.204084846, + "including_init_synchronized_e2e_ms_per_call": 0.6424945270000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.630088, + "synchronized_e2e_ms": 7.659112, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 127.920711, + "median": 0.054496, + "min": 0.042177, + "p90": 0.0791681, + "sample_count": 8290 + }, + "host_enqueue_ms": { + "max": 138.77448, + "median": 0.15968, + "min": 0.128608, + "p90": 0.21863370000000007, + "sample_count": 8290 + }, + "sample_count": 8290, + "synchronized_e2e_ms": { + "max": 149.135067, + "median": 0.179904, + "min": 0.144832, + "p90": 0.2443913, + "sample_count": 8290 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 14502, + "candidate_precomputed_gpu_span_ms": 0.014368, + "candidate_precomputed_host_enqueue_ms": 0.058272, + "candidate_precomputed_inter_kernel_gap_ms": 0.006752, + "candidate_precomputed_kernel_sum_ms": 0.007648, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.058272, + "candidate_precomputed_synchronized_e2e_ms": 0.070816, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.007648 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.063424, + "synchronized_e2e_ms": 0.081248 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.014368 + }, + "host_enqueue_ms": { + "median": 0.058272 + }, + "inter_kernel_gap_ms": { + "median": 0.006752 + }, + "kernel_sum_ms": { + "median": 0.007648 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 14502, + "submission_ms": { + "median": 0.058272 + }, + "synchronized_e2e_ms": { + "median": 0.070816 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc028af890", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc028afce0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 0.846325167037862, + "submission": 0.8638110928061505, + "synchronized_e2e": 0.9306442046995029 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 8290, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.01216, + "candidate_public_raw_host_enqueue_ms": 0.050336, + "candidate_public_raw_inter_kernel_gap_ms": 0.000128, + "candidate_public_raw_kernel_sum_ms": 0.012032, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.050336, + "candidate_public_raw_synchronized_e2e_ms": 0.0659045, + "candidate_public_raw_tflops_from_gpu_span": 9.313010526315788, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.012032 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.061632, + "synchronized_e2e_ms": 0.079456 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.01216 + }, + "host_enqueue_ms": { + "median": 0.050336 + }, + "inter_kernel_gap_ms": { + "median": 0.000128 + }, + "kernel_sum_ms": { + "median": 0.012032 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 8290, + "submission_ms": { + "median": 0.050336 + }, + "synchronized_e2e_ms": { + "median": 0.0659045 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.732835, + "after_init_synchronized_e2e_ms_per_call": 2.759363, + "including_init_host_enqueue_ms_per_call": 40.090858, + "including_init_synchronized_e2e_ms_per_call": 40.194986, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.31858590000000003, + "after_init_synchronized_e2e_ms_per_call": 0.33525035000000003, + "including_init_host_enqueue_ms_per_call": 4.0543882, + "including_init_synchronized_e2e_ms_per_call": 4.07881265, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07716099, + "after_init_synchronized_e2e_ms_per_call": 0.092839085, + "including_init_host_enqueue_ms_per_call": 0.45074121999999994, + "including_init_synchronized_e2e_ms_per_call": 0.467195315, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.053018499, + "after_init_synchronized_e2e_ms_per_call": 0.0685979585, + "including_init_host_enqueue_ms_per_call": 0.090376522, + "including_init_synchronized_e2e_ms_per_call": 0.10603358150000002, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.732835, + "synchronized_e2e_ms": 2.759363, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.012704, + "median": 0.01216, + "min": 0.011808, + "p90": 0.012352, + "sample_count": 8290 + }, + "host_enqueue_ms": { + "max": 91.869119, + "median": 0.050336, + "min": 0.037856, + "p90": 0.07600640000000002, + "sample_count": 8290 + }, + "sample_count": 8290, + "synchronized_e2e_ms": { + "max": 92.091487, + "median": 0.0659045, + "min": 0.053792, + "p90": 0.09693120000000001, + "sample_count": 8290 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.071712, + "submission_ms": 0.071712, + "synchronized_e2e_ms": 0.108288 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.052736, + "submission_ms": 0.052736, + "synchronized_e2e_ms": 0.069888 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.236576, + "submission_ms": 0.236576, + "synchronized_e2e_ms": 0.257312 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.630088, + "submission_ms": 7.630088, + "synchronized_e2e_ms": 7.659112 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.082785, + "submission_ms": 0.082785, + "synchronized_e2e_ms": 0.101473 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.063424, + "submission_ms": 0.063424, + "synchronized_e2e_ms": 0.081248 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.258785, + "submission_ms": 1.258785, + "synchronized_e2e_ms": 1.283521 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.120929, + "submission_ms": 1.120929, + "synchronized_e2e_ms": 1.143905 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.061632, + "submission_ms": 0.061632, + "synchronized_e2e_ms": 0.079456 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.732835, + "submission_ms": 2.732835, + "synchronized_e2e_ms": 2.759363 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.281087, + "evolution_kernel_ms": 0.169183, + "evolution_speedup": 1.6614, + "evolution_tflops": 0.6694, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_5600_d48_low_n_boundary_b3_n1536_k256_d48", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 45584, + "measurement_schedule_sha256": "0639d9c60c56d662d35245fbdabc9cb02770e6111586da070f8c4f42b6c23a49", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 14502, + "public_pair_count": 8290, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 14502, + "baseline_public_raw": 8290, + "candidate_precomputed": 14502, + "candidate_public_raw": 8290 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 9118 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:adjacent_5600_d48_low_n_boundary_b3_n1536_k256_d48", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.47884187082405344, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.729768073500292, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.775681198885395, + "including_init_synchronized_e2e_speedup": 11.513138255602328, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.7675580353607385, + "including_init_synchronized_e2e_speedup": 11.385410580208926, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.743414371220914, + "including_init_synchronized_e2e_speedup": 10.286509979236415, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.731614935741856, + "including_init_synchronized_e2e_speedup": 6.059349480711448, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 4.481578947368421, + "hot_synchronized_e2e_speedup": 2.729768073500292, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 5600481, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_5600_d48_low_n_boundary_b3_n1536_k256_d48", + "source": "guard_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 64, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 5, + "D": 112, + "K": 512, + "N": 2944, + "baseline_07cf_adapter_bench_iters": 4338, + "baseline_07cf_adapter_gpu_span_ms": 0.056512, + "baseline_07cf_adapter_host_enqueue_ms": 0.144032, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.031536499999999995, + "baseline_07cf_adapter_kernel_sum_ms": 0.02496, + "baseline_07cf_adapter_submission_ms": 0.144032, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.1626725, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.02496 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.362528, + "synchronized_e2e_ms": 0.397152 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.056512 + }, + "host_enqueue_ms": { + "median": 0.144032 + }, + "inter_kernel_gap_ms": { + "median": 0.031536499999999995 + }, + "kernel_sum_ms": { + "median": 0.02496 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 4338, + "submission_ms": { + "median": 0.144032 + }, + "synchronized_e2e_ms": { + "median": 0.1626725 + } + }, + "baseline_07cf_precomputed_bench_iters": 6869, + "baseline_07cf_precomputed_gpu_span_ms": 0.014624, + "baseline_07cf_precomputed_host_enqueue_ms": 0.041696, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.014624, + "baseline_07cf_precomputed_submission_ms": 0.041696, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.063008, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.014624 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.079777, + "synchronized_e2e_ms": 0.110049 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.014624 + }, + "host_enqueue_ms": { + "median": 0.041696 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.014624 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 6869, + "submission_ms": { + "median": 0.041696 + }, + "synchronized_e2e_ms": { + "median": 0.063008 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 3.8643326039387307, + "submission": 3.4543361473522642, + "synchronized_e2e": 2.5817753301168107 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 9.102282, + "after_init_synchronized_e2e_ms_per_call": 9.147786, + "including_init_host_enqueue_ms_per_call": 43.330189000000004, + "including_init_synchronized_e2e_ms_per_call": 43.493229, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 1.039857, + "after_init_synchronized_e2e_ms_per_call": 1.06118385, + "including_init_host_enqueue_ms_per_call": 4.4626477, + "including_init_synchronized_e2e_ms_per_call": 4.49572815, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23361449999999997, + "after_init_synchronized_e2e_ms_per_call": 0.252523635, + "including_init_host_enqueue_ms_per_call": 0.5758935700000001, + "including_init_synchronized_e2e_ms_per_call": 0.5959780650000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15299025, + "after_init_synchronized_e2e_ms_per_call": 0.1716576135, + "including_init_host_enqueue_ms_per_call": 0.18721815700000002, + "including_init_synchronized_e2e_ms_per_call": 0.2060030565, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 9.102282, + "synchronized_e2e_ms": 9.147786, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.396416, + "median": 0.056512, + "min": 0.047904, + "p90": 0.08107520000000001, + "sample_count": 4338 + }, + "host_enqueue_ms": { + "max": 56.881275, + "median": 0.144032, + "min": 0.122368, + "p90": 0.2284992, + "sample_count": 4338 + }, + "sample_count": 4338, + "synchronized_e2e_ms": { + "max": 57.069147, + "median": 0.1626725, + "min": 0.137697, + "p90": 0.252096, + "sample_count": 4338 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 6869, + "candidate_precomputed_gpu_span_ms": 0.014944, + "candidate_precomputed_host_enqueue_ms": 0.053344, + "candidate_precomputed_inter_kernel_gap_ms": 0.003776, + "candidate_precomputed_kernel_sum_ms": 0.011168, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.053344, + "candidate_precomputed_synchronized_e2e_ms": 0.065312, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.011168 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.083072, + "synchronized_e2e_ms": 0.110528 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.014944 + }, + "host_enqueue_ms": { + "median": 0.053344 + }, + "inter_kernel_gap_ms": { + "median": 0.003776 + }, + "kernel_sum_ms": { + "median": 0.011168 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 6869, + "submission_ms": { + "median": 0.053344 + }, + "synchronized_e2e_ms": { + "median": 0.065312 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc0273c320", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc0273dee0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.5653104925053531, + "submission": 0.826634673065387, + "synchronized_e2e": 1.0661440470357668 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 4338, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.023392, + "candidate_public_raw_host_enqueue_ms": 0.044096, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.023392, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.044096, + "candidate_public_raw_synchronized_e2e_ms": 0.069632, + "candidate_public_raw_tflops_from_gpu_span": 72.17028727770179, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.023328 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.096128, + "synchronized_e2e_ms": 0.12832 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.023392 + }, + "host_enqueue_ms": { + "median": 0.044096 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.023392 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 4338, + "submission_ms": { + "median": 0.044096 + }, + "synchronized_e2e_ms": { + "median": 0.069632 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.210147, + "after_init_synchronized_e2e_ms_per_call": 3.252451, + "including_init_host_enqueue_ms_per_call": 37.787687, + "including_init_synchronized_e2e_ms_per_call": 451.340243, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.3607011, + "after_init_synchronized_e2e_ms_per_call": 0.38791390000000003, + "including_init_host_enqueue_ms_per_call": 3.8184551, + "including_init_synchronized_e2e_ms_per_call": 45.1966931, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07575651, + "after_init_synchronized_e2e_ms_per_call": 0.10146019, + "including_init_host_enqueue_ms_per_call": 0.42153191, + "including_init_synchronized_e2e_ms_per_call": 4.58233811, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.047262051, + "after_init_synchronized_e2e_ms_per_call": 0.072814819, + "including_init_host_enqueue_ms_per_call": 0.081839591, + "including_init_synchronized_e2e_ms_per_call": 0.520902611, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.210147, + "synchronized_e2e_ms": 3.252451, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.024032, + "median": 0.023392, + "min": 0.022656, + "p90": 0.023584, + "sample_count": 4338 + }, + "host_enqueue_ms": { + "max": 59.067261, + "median": 0.044096, + "min": 0.034336, + "p90": 0.07038720000000001, + "sample_count": 4338 + }, + "sample_count": 4338, + "synchronized_e2e_ms": { + "max": 59.293566, + "median": 0.069632, + "min": 0.060608, + "p90": 0.0924256, + "sample_count": 4338 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.108512, + "submission_ms": 0.108512, + "synchronized_e2e_ms": 0.141536 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.079777, + "submission_ms": 0.079777, + "synchronized_e2e_ms": 0.110049 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.362528, + "submission_ms": 0.362528, + "synchronized_e2e_ms": 0.397152 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 9.102282, + "submission_ms": 9.102282, + "synchronized_e2e_ms": 9.147786 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.123424, + "submission_ms": 0.123424, + "synchronized_e2e_ms": 0.153312 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.083072, + "submission_ms": 0.083072, + "synchronized_e2e_ms": 0.110528 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.729986, + "submission_ms": 1.729986, + "synchronized_e2e_ms": 1.76893 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.637634, + "submission_ms": 1.637634, + "synchronized_e2e_ms": 1.675074 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.096128, + "submission_ms": 0.096128, + "synchronized_e2e_ms": 0.12832 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.210147, + "submission_ms": 3.210147, + "synchronized_e2e_ms": 3.252451 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.296095, + "evolution_kernel_ms": 0.173023, + "evolution_speedup": 1.7113, + "evolution_tflops": 9.7571, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_68cf_d112_tail_b5_n2944_k512_d112", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 22414, + "measurement_schedule_sha256": "b06070a3ff37f0f7904f71eee472ca20fd1033821426685a529edc9988091f05", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 6869, + "public_pair_count": 4338, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 6869, + "baseline_public_raw": 4338, + "candidate_precomputed": 6869, + "candidate_public_raw": 4338 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 4484 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:adjacent_68cf_d112_tail_b5_n2944_k512_d112", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.9785867237687366, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.3361744600183822, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.8125822648826992, + "including_init_synchronized_e2e_speedup": 0.0963646155523517, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.7356169758289144, + "including_init_synchronized_e2e_speedup": 0.09947028956417167, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.4888937720301922, + "including_init_synchronized_e2e_speedup": 0.13005981895997631, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.357454373401656, + "including_init_synchronized_e2e_speedup": 0.3954732653470996, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.415868673050616, + "hot_synchronized_e2e_speedup": 2.3361744600183822, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 6811201, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_68cf_d112_tail_b5_n2944_k512_d112", + "source": "tail_divisibility", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 7, + "D": 128, + "K": 512, + "N": 6016, + "baseline_07cf_adapter_bench_iters": 2259, + "baseline_07cf_adapter_gpu_span_ms": 0.068128, + "baseline_07cf_adapter_host_enqueue_ms": 0.156896, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.02, + "baseline_07cf_adapter_kernel_sum_ms": 0.048128, + "baseline_07cf_adapter_submission_ms": 0.156896, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.177216, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.048128 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.196, + "synchronized_e2e_ms": 0.215584 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.068128 + }, + "host_enqueue_ms": { + "median": 0.156896 + }, + "inter_kernel_gap_ms": { + "median": 0.02 + }, + "kernel_sum_ms": { + "median": 0.048128 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2259, + "submission_ms": { + "median": 0.156896 + }, + "synchronized_e2e_ms": { + "median": 0.177216 + } + }, + "baseline_07cf_precomputed_bench_iters": 5591, + "baseline_07cf_precomputed_gpu_span_ms": 0.02096, + "baseline_07cf_precomputed_host_enqueue_ms": 0.043744, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.02096, + "baseline_07cf_precomputed_submission_ms": 0.043744, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.07024, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.02096 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.042464, + "synchronized_e2e_ms": 0.066976 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.02096 + }, + "host_enqueue_ms": { + "median": 0.043744 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.02096 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5591, + "submission_ms": { + "median": 0.043744 + }, + "synchronized_e2e_ms": { + "median": 0.07024 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 3.2503816793893128, + "submission": 3.5866861741038774, + "synchronized_e2e": 2.5230068337129845 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.155079, + "after_init_synchronized_e2e_ms_per_call": 7.180775, + "including_init_host_enqueue_ms_per_call": 41.604523, + "including_init_synchronized_e2e_ms_per_call": 491.222269, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8567142999999999, + "after_init_synchronized_e2e_ms_per_call": 0.8775719000000001, + "including_init_host_enqueue_ms_per_call": 4.3016587, + "including_init_synchronized_e2e_ms_per_call": 49.2817213, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22687783, + "after_init_synchronized_e2e_ms_per_call": 0.24725159000000002, + "including_init_host_enqueue_ms_per_call": 0.57137227, + "including_init_synchronized_e2e_ms_per_call": 5.08766653, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.163894183, + "after_init_synchronized_e2e_ms_per_call": 0.18421955900000003, + "including_init_host_enqueue_ms_per_call": 0.198343627, + "including_init_synchronized_e2e_ms_per_call": 0.6682610529999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.155079, + "synchronized_e2e_ms": 7.180775, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.203045, + "median": 0.068128, + "min": 0.056576, + "p90": 0.08671360000000002, + "sample_count": 2259 + }, + "host_enqueue_ms": { + "max": 30.892544, + "median": 0.156896, + "min": 0.12912, + "p90": 0.1965698, + "sample_count": 2259 + }, + "sample_count": 2259, + "synchronized_e2e_ms": { + "max": 30.931328, + "median": 0.177216, + "min": 0.147136, + "p90": 0.22048020000000002, + "sample_count": 2259 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 5591, + "candidate_precomputed_gpu_span_ms": 0.018144, + "candidate_precomputed_host_enqueue_ms": 0.04112, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.018144, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.04112, + "candidate_precomputed_synchronized_e2e_ms": 0.053312, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.018144 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.031168, + "synchronized_e2e_ms": 0.044928 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.018144 + }, + "host_enqueue_ms": { + "median": 0.04112 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.018144 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5591, + "submission_ms": { + "median": 0.04112 + }, + "synchronized_e2e_ms": { + "median": 0.053312 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc045eb500", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc045ea150" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.46031746031746, + "submission": 1.1992217898832687, + "synchronized_e2e": 1.7454981992797118 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 2259, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.04464, + "candidate_public_raw_host_enqueue_ms": 0.049312, + "candidate_public_raw_inter_kernel_gap_ms": 0.00016, + "candidate_public_raw_kernel_sum_ms": 0.04448, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.049312, + "candidate_public_raw_synchronized_e2e_ms": 0.093056, + "candidate_public_raw_tflops_from_gpu_span": 123.64928458781363, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.04448 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.271104, + "synchronized_e2e_ms": 0.315585 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.04464 + }, + "host_enqueue_ms": { + "median": 0.049312 + }, + "inter_kernel_gap_ms": { + "median": 0.00016 + }, + "kernel_sum_ms": { + "median": 0.04448 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2259, + "submission_ms": { + "median": 0.049312 + }, + "synchronized_e2e_ms": { + "median": 0.093056 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.243075, + "after_init_synchronized_e2e_ms_per_call": 2.271043, + "including_init_host_enqueue_ms_per_call": 37.117158999999994, + "including_init_synchronized_e2e_ms_per_call": 37.214119, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.26868830000000005, + "after_init_synchronized_e2e_ms_per_call": 0.31085470000000004, + "including_init_host_enqueue_ms_per_call": 3.756096699999999, + "including_init_synchronized_e2e_ms_per_call": 3.8051623, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07124963000000001, + "after_init_synchronized_e2e_ms_per_call": 0.11483587, + "including_init_host_enqueue_ms_per_call": 0.4199904699999999, + "including_init_synchronized_e2e_ms_per_call": 0.46426662999999996, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.051505763, + "after_init_synchronized_e2e_ms_per_call": 0.09523398699999999, + "including_init_host_enqueue_ms_per_call": 0.086379847, + "including_init_synchronized_e2e_ms_per_call": 0.13017706299999998, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.243075, + "synchronized_e2e_ms": 2.271043, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.045536, + "median": 0.04464, + "min": 0.044064, + "p90": 0.044992, + "sample_count": 2259 + }, + "host_enqueue_ms": { + "max": 15.404656, + "median": 0.049312, + "min": 0.03776, + "p90": 0.0685824, + "sample_count": 2259 + }, + "sample_count": 2259, + "synchronized_e2e_ms": { + "max": 20.53999, + "median": 0.093056, + "min": 0.083936, + "p90": 0.1103296, + "sample_count": 2259 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.063904, + "submission_ms": 0.063904, + "synchronized_e2e_ms": 0.087808 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.042464, + "submission_ms": 0.042464, + "synchronized_e2e_ms": 0.066976 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.196, + "submission_ms": 0.196, + "synchronized_e2e_ms": 0.215584 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.155079, + "submission_ms": 7.155079, + "synchronized_e2e_ms": 7.180775 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.05168, + "submission_ms": 0.05168, + "synchronized_e2e_ms": 0.067264 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.031168, + "submission_ms": 0.031168, + "synchronized_e2e_ms": 0.044928 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.184609, + "submission_ms": 1.184609, + "synchronized_e2e_ms": 1.206689 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.012865, + "submission_ms": 1.012865, + "synchronized_e2e_ms": 1.032737 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.271104, + "submission_ms": 0.271104, + "synchronized_e2e_ms": 0.315585 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.243075, + "submission_ms": 2.243075, + "synchronized_e2e_ms": 2.271043 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.386527, + "evolution_kernel_ms": 0.156832, + "evolution_speedup": 2.4646, + "evolution_tflops": 35.195, + "expected_route": "aligned_weave_v10_fallback", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_68cf_d128_forced_fallback_b7_n6016_k512_d128", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 15700, + "measurement_schedule_sha256": "6997e54a04ac15dd8d1861d6c29b49034d6e4838f077c8572f9f1becbfbd6dc5", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 5591, + "public_pair_count": 2259, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 5591, + "baseline_public_raw": 2259, + "candidate_precomputed": 5591, + "candidate_public_raw": 2259 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 3142 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:adjacent_68cf_d128_forced_fallback_b7_n6016_k512_d128", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.1552028218694885, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.9044016506189823, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.1618842091497164, + "including_init_synchronized_e2e_speedup": 13.199889778392981, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.823093554641445, + "including_init_synchronized_e2e_speedup": 12.951279712825915, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.153086748940031, + "including_init_synchronized_e2e_speedup": 10.95850143267889, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.9343888122630006, + "including_init_synchronized_e2e_speedup": 5.133477723337482, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.5261648745519711, + "hot_synchronized_e2e_speedup": 1.9044016506189823, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 6812801, + "selected_route": "aligned_weave_v10_fallback", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_68cf_d128_forced_fallback_b7_n6016_k512_d128", + "source": "forced_fallback", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 3, + "D": 224, + "K": 1024, + "N": 5120, + "baseline_07cf_adapter_bench_iters": 2465, + "baseline_07cf_adapter_gpu_span_ms": 0.075265, + "baseline_07cf_adapter_host_enqueue_ms": 0.147264, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.021312, + "baseline_07cf_adapter_kernel_sum_ms": 0.053952, + "baseline_07cf_adapter_submission_ms": 0.147264, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.174144, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.053952 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.186593, + "synchronized_e2e_ms": 0.211937 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.075265 + }, + "host_enqueue_ms": { + "median": 0.147264 + }, + "inter_kernel_gap_ms": { + "median": 0.021312 + }, + "kernel_sum_ms": { + "median": 0.053952 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2465, + "submission_ms": { + "median": 0.147264 + }, + "synchronized_e2e_ms": { + "median": 0.174144 + } + }, + "baseline_07cf_precomputed_bench_iters": 4178, + "baseline_07cf_precomputed_gpu_span_ms": 0.033088, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042688, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.033088, + "baseline_07cf_precomputed_submission_ms": 0.042688, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.081184, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.033088 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.040288, + "synchronized_e2e_ms": 0.076 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.033088 + }, + "host_enqueue_ms": { + "median": 0.042688 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.033088 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4178, + "submission_ms": { + "median": 0.042688 + }, + "synchronized_e2e_ms": { + "median": 0.081184 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.2746917311411994, + "submission": 3.4497751124437785, + "synchronized_e2e": 2.145053212455656 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.027783, + "after_init_synchronized_e2e_ms_per_call": 7.051943, + "including_init_host_enqueue_ms_per_call": 42.579051, + "including_init_synchronized_e2e_ms_per_call": 42.68318, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8353159, + "after_init_synchronized_e2e_ms_per_call": 0.8619239000000001, + "including_init_host_enqueue_ms_per_call": 4.3904426999999995, + "including_init_synchronized_e2e_ms_per_call": 4.4250476, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.21606919000000002, + "after_init_synchronized_e2e_ms_per_call": 0.24292198999999998, + "including_init_host_enqueue_ms_per_call": 0.5715818699999999, + "including_init_synchronized_e2e_ms_per_call": 0.59923436, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.154144519, + "after_init_synchronized_e2e_ms_per_call": 0.18102179899999998, + "including_init_host_enqueue_ms_per_call": 0.189695787, + "including_init_synchronized_e2e_ms_per_call": 0.216653036, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.027783, + "synchronized_e2e_ms": 7.051943, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 46.141522, + "median": 0.075265, + "min": 0.066592, + "p90": 0.0966726, + "sample_count": 2465 + }, + "host_enqueue_ms": { + "max": 50.50034, + "median": 0.147264, + "min": 0.126688, + "p90": 0.23621159999999997, + "sample_count": 2465 + }, + "sample_count": 2465, + "synchronized_e2e_ms": { + "max": 55.660986, + "median": 0.174144, + "min": 0.155104, + "p90": 0.2608579999999999, + "sample_count": 2465 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 4178, + "candidate_precomputed_gpu_span_ms": 0.020416, + "candidate_precomputed_host_enqueue_ms": 0.053216, + "candidate_precomputed_inter_kernel_gap_ms": 0.002304, + "candidate_precomputed_kernel_sum_ms": 0.018144, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.053216, + "candidate_precomputed_synchronized_e2e_ms": 0.06496, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.018144 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.040896, + "synchronized_e2e_ms": 0.054112 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.020416 + }, + "host_enqueue_ms": { + "median": 0.053216 + }, + "inter_kernel_gap_ms": { + "median": 0.002304 + }, + "kernel_sum_ms": { + "median": 0.018144 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4178, + "submission_ms": { + "median": 0.053216 + }, + "synchronized_e2e_ms": { + "median": 0.06496 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc02715e50", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc02714530" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.9858934169278994, + "submission": 0.8394467829224295, + "synchronized_e2e": 1.3275862068965516 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2465, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.040544, + "candidate_public_raw_host_enqueue_ms": 0.044672, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.040768, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.044672, + "candidate_public_raw_synchronized_e2e_ms": 0.08624, + "candidate_public_raw_tflops_from_gpu_span": 173.7971270718232, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.040544 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.051072, + "synchronized_e2e_ms": 0.08032 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.040544 + }, + "host_enqueue_ms": { + "median": 0.044672 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.040768 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2465, + "submission_ms": { + "median": 0.044672 + }, + "synchronized_e2e_ms": { + "median": 0.08624 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.14077, + "after_init_synchronized_e2e_ms_per_call": 2.16589, + "including_init_host_enqueue_ms_per_call": 37.978790000000004, + "including_init_synchronized_e2e_ms_per_call": 486.266103, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2542818, + "after_init_synchronized_e2e_ms_per_call": 0.294205, + "including_init_host_enqueue_ms_per_call": 3.8380838000000006, + "including_init_synchronized_e2e_ms_per_call": 48.7042263, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06563298000000001, + "after_init_synchronized_e2e_ms_per_call": 0.10703649999999999, + "including_init_host_enqueue_ms_per_call": 0.42401318000000005, + "including_init_synchronized_e2e_ms_per_call": 4.94803863, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.04676809800000001, + "after_init_synchronized_e2e_ms_per_call": 0.08831965, + "including_init_host_enqueue_ms_per_call": 0.082606118, + "including_init_synchronized_e2e_ms_per_call": 0.5724198629999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.14077, + "synchronized_e2e_ms": 2.16589, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.041248, + "median": 0.040544, + "min": 0.040064, + "p90": 0.040832, + "sample_count": 2465 + }, + "host_enqueue_ms": { + "max": 9.26817, + "median": 0.044672, + "min": 0.034848, + "p90": 0.0723076, + "sample_count": 2465 + }, + "sample_count": 2465, + "synchronized_e2e_ms": { + "max": 9.297482, + "median": 0.08624, + "min": 0.07936, + "p90": 0.10742400000000002, + "sample_count": 2465 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.06512, + "submission_ms": 0.06512, + "synchronized_e2e_ms": 0.09904 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.040288, + "submission_ms": 0.040288, + "synchronized_e2e_ms": 0.076 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.186593, + "submission_ms": 0.186593, + "synchronized_e2e_ms": 0.211937 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.027783, + "submission_ms": 7.027783, + "synchronized_e2e_ms": 7.051943 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.092896, + "submission_ms": 0.092896, + "synchronized_e2e_ms": 0.110368 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.040896, + "submission_ms": 0.040896, + "synchronized_e2e_ms": 0.054112 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.126689, + "submission_ms": 1.126689, + "synchronized_e2e_ms": 1.151041 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.057729, + "submission_ms": 1.057729, + "synchronized_e2e_ms": 1.080705 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.051072, + "submission_ms": 0.051072, + "synchronized_e2e_ms": 0.08032 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.14077, + "submission_ms": 2.14077, + "synchronized_e2e_ms": 2.16589 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.486208, + "evolution_kernel_ms": 0.180032, + "evolution_speedup": 2.7007, + "evolution_tflops": 39.1399, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_68cf_d224_overlap_b3_n5120_k1024_d224", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 13286, + "measurement_schedule_sha256": "cf92a346cc7cc13ab33d7d3e39341a2e154e635a5ac90e786e2509f213dd7216", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 4178, + "public_pair_count": 2465, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 4178, + "baseline_public_raw": 2465, + "candidate_precomputed": 4178, + "candidate_public_raw": 2465 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2658 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:adjacent_68cf_d224_overlap_b3_n5120_k1024_d224", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.6206896551724137, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.019294990723562, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.255910041599527, + "including_init_synchronized_e2e_speedup": 0.08777741186701636, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.9296711476691426, + "including_init_synchronized_e2e_speedup": 0.09085551575634002, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.2695247882731593, + "including_init_synchronized_e2e_speedup": 0.12110543284097197, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.0496208827820306, + "including_init_synchronized_e2e_speedup": 0.3784862301327933, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.8563782557221784, + "hot_synchronized_e2e_speedup": 2.019294990723562, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 6822401, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_68cf_d224_overlap_b3_n5120_k1024_d224", + "source": "guard_overlap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 288, + "K": 512, + "N": 1920, + "baseline_07cf_adapter_bench_iters": 3701, + "baseline_07cf_adapter_gpu_span_ms": 0.063872, + "baseline_07cf_adapter_host_enqueue_ms": 0.155424, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.038913, + "baseline_07cf_adapter_kernel_sum_ms": 0.02496, + "baseline_07cf_adapter_submission_ms": 0.155424, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.174784, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.02496 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.192449, + "synchronized_e2e_ms": 0.209089 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.063872 + }, + "host_enqueue_ms": { + "median": 0.155424 + }, + "inter_kernel_gap_ms": { + "median": 0.038913 + }, + "kernel_sum_ms": { + "median": 0.02496 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3701, + "submission_ms": { + "median": 0.155424 + }, + "synchronized_e2e_ms": { + "median": 0.174784 + } + }, + "baseline_07cf_precomputed_bench_iters": 4647, + "baseline_07cf_precomputed_gpu_span_ms": 0.021952, + "baseline_07cf_precomputed_host_enqueue_ms": 0.046464, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.021952, + "baseline_07cf_precomputed_submission_ms": 0.046464, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.075488, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.021952 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.041152, + "synchronized_e2e_ms": 0.062688 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.021952 + }, + "host_enqueue_ms": { + "median": 0.046464 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.021952 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4647, + "submission_ms": { + "median": 0.046464 + }, + "synchronized_e2e_ms": { + "median": 0.075488 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.9096209912536444, + "submission": 3.34504132231405, + "synchronized_e2e": 2.3153878762187365 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.95924, + "after_init_synchronized_e2e_ms_per_call": 7.985544, + "including_init_host_enqueue_ms_per_call": 44.893678, + "including_init_synchronized_e2e_ms_per_call": 463.09686300000004, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9358056000000001, + "after_init_synchronized_e2e_ms_per_call": 0.95586, + "including_init_host_enqueue_ms_per_call": 4.629249400000001, + "including_init_synchronized_e2e_ms_per_call": 46.466991900000004, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23346216000000003, + "after_init_synchronized_e2e_ms_per_call": 0.2528916, + "including_init_host_enqueue_ms_per_call": 0.60280654, + "including_init_synchronized_e2e_ms_per_call": 4.8040047900000005, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.163227816, + "after_init_synchronized_e2e_ms_per_call": 0.18259476, + "including_init_host_enqueue_ms_per_call": 0.20016225399999998, + "including_init_synchronized_e2e_ms_per_call": 0.6377060790000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.95924, + "synchronized_e2e_ms": 7.985544, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 75.734478, + "median": 0.063872, + "min": 0.053696, + "p90": 0.089248, + "sample_count": 3701 + }, + "host_enqueue_ms": { + "max": 75.876238, + "median": 0.155424, + "min": 0.131136, + "p90": 0.221536, + "sample_count": 3701 + }, + "sample_count": 3701, + "synchronized_e2e_ms": { + "max": 75.940174, + "median": 0.174784, + "min": 0.149024, + "p90": 0.247136, + "sample_count": 3701 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 4647, + "candidate_precomputed_gpu_span_ms": 0.021568, + "candidate_precomputed_host_enqueue_ms": 0.043232, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.021568, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.043232, + "candidate_precomputed_synchronized_e2e_ms": 0.058848, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.021568 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.031328, + "synchronized_e2e_ms": 0.044768 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.021568 + }, + "host_enqueue_ms": { + "median": 0.043232 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.021568 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4647, + "submission_ms": { + "median": 0.043232 + }, + "synchronized_e2e_ms": { + "median": 0.058848 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01aee7b0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01aeeab0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.2537091988130564, + "submission": 1.1228719467061434, + "synchronized_e2e": 1.3039695486677543 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 3701, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.02704, + "candidate_public_raw_host_enqueue_ms": 0.048544, + "candidate_public_raw_inter_kernel_gap_ms": 0.000192, + "candidate_public_raw_kernel_sum_ms": 0.026848, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.048544, + "candidate_public_raw_synchronized_e2e_ms": 0.076736, + "candidate_public_raw_tflops_from_gpu_span": 41.88099408284023, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.026848 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.236352, + "synchronized_e2e_ms": 0.255968 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.02704 + }, + "host_enqueue_ms": { + "median": 0.048544 + }, + "inter_kernel_gap_ms": { + "median": 0.000192 + }, + "kernel_sum_ms": { + "median": 0.026848 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3701, + "submission_ms": { + "median": 0.048544 + }, + "synchronized_e2e_ms": { + "median": 0.076736 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.124579, + "after_init_synchronized_e2e_ms_per_call": 3.151139, + "including_init_host_enqueue_ms_per_call": 40.48260199999999, + "including_init_synchronized_e2e_ms_per_call": 40.586762, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.3561475, + "after_init_synchronized_e2e_ms_per_call": 0.3841763, + "including_init_host_enqueue_ms_per_call": 4.091949799999999, + "including_init_synchronized_e2e_ms_per_call": 4.1277386, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07930435, + "after_init_synchronized_e2e_ms_per_call": 0.10748003, + "including_init_host_enqueue_ms_per_call": 0.4528845799999999, + "including_init_synchronized_e2e_ms_per_call": 0.48183626, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.051620034999999995, + "after_init_synchronized_e2e_ms_per_call": 0.07981040299999999, + "including_init_host_enqueue_ms_per_call": 0.08897805799999998, + "including_init_synchronized_e2e_ms_per_call": 0.117246026, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.124579, + "synchronized_e2e_ms": 3.151139, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.027744, + "median": 0.02704, + "min": 0.026592, + "p90": 0.027296, + "sample_count": 3701 + }, + "host_enqueue_ms": { + "max": 75.673646, + "median": 0.048544, + "min": 0.037152, + "p90": 0.075552, + "sample_count": 3701 + }, + "sample_count": 3701, + "synchronized_e2e_ms": { + "max": 80.821075, + "median": 0.076736, + "min": 0.066688, + "p90": 0.101312, + "sample_count": 3701 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.067553, + "submission_ms": 0.067553, + "synchronized_e2e_ms": 0.087617 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.041152, + "submission_ms": 0.041152, + "synchronized_e2e_ms": 0.062688 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.192449, + "submission_ms": 0.192449, + "synchronized_e2e_ms": 0.209089 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.95924, + "submission_ms": 7.95924, + "synchronized_e2e_ms": 7.985544 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.079712, + "submission_ms": 0.079712, + "synchronized_e2e_ms": 0.09488 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.031328, + "submission_ms": 0.031328, + "synchronized_e2e_ms": 0.044768 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.229281, + "submission_ms": 1.229281, + "synchronized_e2e_ms": 1.249569 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.056385, + "submission_ms": 1.056385, + "synchronized_e2e_ms": 1.076737 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.236352, + "submission_ms": 0.236352, + "synchronized_e2e_ms": 0.255968 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.124579, + "submission_ms": 3.124579, + "synchronized_e2e_ms": 3.151139 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.421007, + "evolution_kernel_ms": 0.177824, + "evolution_speedup": 2.3676, + "evolution_tflops": 6.3684, + "expected_route": "d288_parent_splitk_hybrid_20260629_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_68cf_d288_boundary_b2_n1920_k512_d288", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 16696, + "measurement_schedule_sha256": "22c54ef8183a937dc4361ca9ebc020fbfcc4ea2429a2cb1ed901718c04f827c9", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 4647, + "public_pair_count": 3701, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 4647, + "baseline_public_raw": 3701, + "candidate_precomputed": 4647, + "candidate_public_raw": 3701 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 3342 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:adjacent_68cf_d288_boundary_b2_n1920_k512_d288", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.0178041543026706, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.2777314428690576, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.5341770071075884, + "including_init_synchronized_e2e_speedup": 11.410047024692437, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.488076437822947, + "including_init_synchronized_e2e_speedup": 11.257251585650314, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.3529170953897203, + "including_init_synchronized_e2e_speedup": 9.9702018897457, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.287856634428973, + "including_init_synchronized_e2e_speedup": 5.439042164209472, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.3621301775147927, + "hot_synchronized_e2e_speedup": 2.2777314428690576, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 6828801, + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_68cf_d288_boundary_b2_n1920_k512_d288", + "source": "guard_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 3, + "D": 352, + "K": 768, + "N": 2816, + "baseline_07cf_adapter_bench_iters": 2333, + "baseline_07cf_adapter_gpu_span_ms": 0.068608, + "baseline_07cf_adapter_host_enqueue_ms": 0.14544, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.025888, + "baseline_07cf_adapter_kernel_sum_ms": 0.04272, + "baseline_07cf_adapter_submission_ms": 0.14544, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.168096, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.04272 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.205568, + "synchronized_e2e_ms": 0.224896 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.068608 + }, + "host_enqueue_ms": { + "median": 0.14544 + }, + "inter_kernel_gap_ms": { + "median": 0.025888 + }, + "kernel_sum_ms": { + "median": 0.04272 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2333, + "submission_ms": { + "median": 0.14544 + }, + "synchronized_e2e_ms": { + "median": 0.168096 + } + }, + "baseline_07cf_precomputed_bench_iters": 3060, + "baseline_07cf_precomputed_gpu_span_ms": 0.033504, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042432, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.033504, + "baseline_07cf_precomputed_submission_ms": 0.042432, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.082528, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.033504 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.040512, + "synchronized_e2e_ms": 0.067201 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.033504 + }, + "host_enqueue_ms": { + "median": 0.042432 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.033504 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3060, + "submission_ms": { + "median": 0.042432 + }, + "synchronized_e2e_ms": { + "median": 0.082528 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.0477554918815666, + "submission": 3.427601809954752, + "synchronized_e2e": 2.0368359829391234 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.321511, + "after_init_synchronized_e2e_ms_per_call": 7.356871, + "including_init_host_enqueue_ms_per_call": 41.549418, + "including_init_synchronized_e2e_ms_per_call": 41.702314, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8630471, + "after_init_synchronized_e2e_ms_per_call": 0.8869735000000001, + "including_init_host_enqueue_ms_per_call": 4.2858378, + "including_init_synchronized_e2e_ms_per_call": 4.3215178000000005, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.21720071000000002, + "after_init_synchronized_e2e_ms_per_call": 0.23998375000000002, + "including_init_host_enqueue_ms_per_call": 0.5594797800000001, + "including_init_synchronized_e2e_ms_per_call": 0.58343818, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.152616071, + "after_init_synchronized_e2e_ms_per_call": 0.175284775, + "including_init_host_enqueue_ms_per_call": 0.18684397800000002, + "including_init_synchronized_e2e_ms_per_call": 0.20963021799999998, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.321511, + "synchronized_e2e_ms": 7.356871, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.345665, + "median": 0.068608, + "min": 0.059776, + "p90": 0.08896640000000003, + "sample_count": 2333 + }, + "host_enqueue_ms": { + "max": 35.609701, + "median": 0.14544, + "min": 0.124544, + "p90": 0.21598080000000003, + "sample_count": 2333 + }, + "sample_count": 2333, + "synchronized_e2e_ms": { + "max": 35.650341, + "median": 0.168096, + "min": 0.149024, + "p90": 0.23960959999999998, + "sample_count": 2333 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3060, + "candidate_precomputed_gpu_span_ms": 0.028096, + "candidate_precomputed_host_enqueue_ms": 0.0543365, + "candidate_precomputed_inter_kernel_gap_ms": 0.002816, + "candidate_precomputed_kernel_sum_ms": 0.02528, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.0543365, + "candidate_precomputed_synchronized_e2e_ms": 0.066272, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.02528 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.039072, + "synchronized_e2e_ms": 0.05488 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.028096 + }, + "host_enqueue_ms": { + "median": 0.0543365 + }, + "inter_kernel_gap_ms": { + "median": 0.002816 + }, + "kernel_sum_ms": { + "median": 0.02528 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3060, + "submission_ms": { + "median": 0.0543365 + }, + "synchronized_e2e_ms": { + "median": 0.066272 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc0270d370", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc0270f4d0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.5330296127562641, + "submission": 0.8085909103457161, + "synchronized_e2e": 1.3471752776436505 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2333, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.043072, + "candidate_public_raw_host_enqueue_ms": 0.043936, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.043264, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.043936, + "candidate_public_raw_synchronized_e2e_ms": 0.08928, + "candidate_public_raw_tflops_from_gpu_span": 106.04562258543834, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.043072 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.070209, + "synchronized_e2e_ms": 0.098625 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.043072 + }, + "host_enqueue_ms": { + "median": 0.043936 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.043264 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2333, + "submission_ms": { + "median": 0.043936 + }, + "synchronized_e2e_ms": { + "median": 0.08928 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.493347, + "after_init_synchronized_e2e_ms_per_call": 2.516867, + "including_init_host_enqueue_ms_per_call": 37.070887, + "including_init_synchronized_e2e_ms_per_call": 450.60465899999997, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2888771, + "after_init_synchronized_e2e_ms_per_call": 0.3320387, + "including_init_host_enqueue_ms_per_call": 3.7466310999999997, + "including_init_synchronized_e2e_ms_per_call": 45.140817899999995, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06843011, + "after_init_synchronized_e2e_ms_per_call": 0.11355587, + "including_init_host_enqueue_ms_per_call": 0.41420551000000005, + "including_init_synchronized_e2e_ms_per_call": 4.59443379, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.046385411, + "after_init_synchronized_e2e_ms_per_call": 0.09170758700000001, + "including_init_host_enqueue_ms_per_call": 0.080962951, + "including_init_synchronized_e2e_ms_per_call": 0.539795379, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.493347, + "synchronized_e2e_ms": 2.516867, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.04368, + "median": 0.043072, + "min": 0.042624, + "p90": 0.043296, + "sample_count": 2333 + }, + "host_enqueue_ms": { + "max": 0.589761, + "median": 0.043936, + "min": 0.034816, + "p90": 0.06486420000000002, + "sample_count": 2333 + }, + "sample_count": 2333, + "synchronized_e2e_ms": { + "max": 0.680001, + "median": 0.08928, + "min": 0.08192, + "p90": 0.1072256, + "sample_count": 2333 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.061472, + "submission_ms": 0.061472, + "synchronized_e2e_ms": 0.087168 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.040512, + "submission_ms": 0.040512, + "synchronized_e2e_ms": 0.067201 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.205568, + "submission_ms": 0.205568, + "synchronized_e2e_ms": 0.224896 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.321511, + "submission_ms": 7.321511, + "synchronized_e2e_ms": 7.356871 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.068096, + "submission_ms": 0.068096, + "synchronized_e2e_ms": 0.083488 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.039072, + "submission_ms": 0.039072, + "synchronized_e2e_ms": 0.05488 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.213697, + "submission_ms": 1.213697, + "synchronized_e2e_ms": 1.236897 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.030017, + "submission_ms": 1.030017, + "synchronized_e2e_ms": 1.051777 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.070209, + "submission_ms": 0.070209, + "synchronized_e2e_ms": 0.098625 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.493347, + "submission_ms": 2.493347, + "synchronized_e2e_ms": 2.516867 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.500207, + "evolution_kernel_ms": 0.18768, + "evolution_speedup": 2.6652, + "evolution_tflops": 24.3372, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_68cf_d352_tail_b3_n2816_k768_d352", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 10786, + "measurement_schedule_sha256": "697b9daed1250a0acf9251442a4757082b158d273bd64fccacd51b3aca1481df", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3060, + "public_pair_count": 2333, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3060, + "baseline_public_raw": 2333, + "candidate_precomputed": 3060, + "candidate_public_raw": 2333 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2158 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:adjacent_68cf_d352_tail_b3_n2816_k768_d352", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.1924829157175398, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.8827956989247312, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.9230273192822662, + "including_init_synchronized_e2e_speedup": 0.092547454108769, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.671295544766318, + "including_init_synchronized_e2e_speedup": 0.09573414929196489, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.1133539816127516, + "including_init_synchronized_e2e_speedup": 0.1269880482922358, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.9113443144022533, + "including_init_synchronized_e2e_speedup": 0.38835126448905744, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.5928677563150075, + "hot_synchronized_e2e_speedup": 1.8827956989247312, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 6835201, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_68cf_d352_tail_b3_n2816_k768_d352", + "source": "tail_divisibility", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 416, + "K": 1024, + "N": 2304, + "baseline_07cf_adapter_bench_iters": 2137, + "baseline_07cf_adapter_gpu_span_ms": 0.079808, + "baseline_07cf_adapter_host_enqueue_ms": 0.160096, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.037056, + "baseline_07cf_adapter_kernel_sum_ms": 0.04272, + "baseline_07cf_adapter_submission_ms": 0.160096, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.186528, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.04272 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.211105, + "synchronized_e2e_ms": 0.233121 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.079808 + }, + "host_enqueue_ms": { + "median": 0.160096 + }, + "inter_kernel_gap_ms": { + "median": 0.037056 + }, + "kernel_sum_ms": { + "median": 0.04272 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2137, + "submission_ms": { + "median": 0.160096 + }, + "synchronized_e2e_ms": { + "median": 0.186528 + } + }, + "baseline_07cf_precomputed_bench_iters": 2844, + "baseline_07cf_precomputed_gpu_span_ms": 0.035168, + "baseline_07cf_precomputed_host_enqueue_ms": 0.047488, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.035168, + "baseline_07cf_precomputed_submission_ms": 0.047488, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.087744, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.035168 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.044608, + "synchronized_e2e_ms": 0.076512 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.035168 + }, + "host_enqueue_ms": { + "median": 0.047488 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.035168 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2844, + "submission_ms": { + "median": 0.047488 + }, + "synchronized_e2e_ms": { + "median": 0.087744 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.2693357597816197, + "submission": 3.371293800539083, + "synchronized_e2e": 2.12582056892779 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.032679, + "after_init_synchronized_e2e_ms_per_call": 7.059015, + "including_init_host_enqueue_ms_per_call": 41.482123, + "including_init_synchronized_e2e_ms_per_call": 491.100509, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8473542999999999, + "after_init_synchronized_e2e_ms_per_call": 0.8737767, + "including_init_host_enqueue_ms_per_call": 4.2922987, + "including_init_synchronized_e2e_ms_per_call": 49.277926099999995, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22882182999999998, + "after_init_synchronized_e2e_ms_per_call": 0.25525286999999997, + "including_init_host_enqueue_ms_per_call": 0.57331627, + "including_init_synchronized_e2e_ms_per_call": 5.09566781, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.166968583, + "after_init_synchronized_e2e_ms_per_call": 0.193400487, + "including_init_host_enqueue_ms_per_call": 0.201418027, + "including_init_synchronized_e2e_ms_per_call": 0.677441981, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.032679, + "synchronized_e2e_ms": 7.059015, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.422209, + "median": 0.079808, + "min": 0.068192, + "p90": 0.09825280000000002, + "sample_count": 2137 + }, + "host_enqueue_ms": { + "max": 34.20138, + "median": 0.160096, + "min": 0.130688, + "p90": 0.1944132, + "sample_count": 2137 + }, + "sample_count": 2137, + "synchronized_e2e_ms": { + "max": 34.237444, + "median": 0.186528, + "min": 0.15856, + "p90": 0.2195264, + "sample_count": 2137 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 2844, + "candidate_precomputed_gpu_span_ms": 0.037984, + "candidate_precomputed_host_enqueue_ms": 0.0569925, + "candidate_precomputed_inter_kernel_gap_ms": 0.003808, + "candidate_precomputed_kernel_sum_ms": 0.034176, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.0569925, + "candidate_precomputed_synchronized_e2e_ms": 0.0745125, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.034176 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.042368, + "synchronized_e2e_ms": 0.064672 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.037984 + }, + "host_enqueue_ms": { + "median": 0.0569925 + }, + "inter_kernel_gap_ms": { + "median": 0.003808 + }, + "kernel_sum_ms": { + "median": 0.034176 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2844, + "submission_ms": { + "median": 0.0569925 + }, + "synchronized_e2e_ms": { + "median": 0.0745125 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295bf0e0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295bc4d0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.2325189553496212, + "submission": 0.9966223625915691, + "synchronized_e2e": 1.388438181513169 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2137, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.046816, + "candidate_public_raw_host_enqueue_ms": 0.0568, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.04672, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.0568, + "candidate_public_raw_synchronized_e2e_ms": 0.103456, + "candidate_public_raw_tflops_from_gpu_span": 83.85741079972658, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.04672 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.07504, + "synchronized_e2e_ms": 0.105216 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.046816 + }, + "host_enqueue_ms": { + "median": 0.0568 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.04672 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2137, + "submission_ms": { + "median": 0.0568 + }, + "synchronized_e2e_ms": { + "median": 0.103456 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.269538, + "after_init_synchronized_e2e_ms_per_call": 2.296898, + "including_init_host_enqueue_ms_per_call": 37.14362199999999, + "including_init_synchronized_e2e_ms_per_call": 37.239974, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2780738, + "after_init_synchronized_e2e_ms_per_call": 0.3228002, + "including_init_host_enqueue_ms_per_call": 3.7654821999999997, + "including_init_synchronized_e2e_ms_per_call": 3.8171077999999996, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07892738, + "after_init_synchronized_e2e_ms_per_call": 0.12539042000000003, + "including_init_host_enqueue_ms_per_call": 0.42766821999999993, + "including_init_synchronized_e2e_ms_per_call": 0.47482118, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.059012738, + "after_init_synchronized_e2e_ms_per_call": 0.10564944200000001, + "including_init_host_enqueue_ms_per_call": 0.093886822, + "including_init_synchronized_e2e_ms_per_call": 0.140592518, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.269538, + "synchronized_e2e_ms": 2.296898, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.047456, + "median": 0.046816, + "min": 0.046336, + "p90": 0.047072, + "sample_count": 2137 + }, + "host_enqueue_ms": { + "max": 44.311886, + "median": 0.0568, + "min": 0.037408, + "p90": 0.073376, + "sample_count": 2137 + }, + "sample_count": 2137, + "synchronized_e2e_ms": { + "max": 44.42859, + "median": 0.103456, + "min": 0.086464, + "p90": 0.11683840000000001, + "sample_count": 2137 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.068, + "submission_ms": 0.068, + "synchronized_e2e_ms": 0.098432 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.044608, + "submission_ms": 0.044608, + "synchronized_e2e_ms": 0.076512 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.211105, + "submission_ms": 0.211105, + "synchronized_e2e_ms": 0.233121 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.032679, + "submission_ms": 7.032679, + "synchronized_e2e_ms": 7.059015 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.072384, + "submission_ms": 0.072384, + "synchronized_e2e_ms": 0.090624 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.042368, + "submission_ms": 0.042368, + "synchronized_e2e_ms": 0.064672 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.230753, + "submission_ms": 1.230753, + "synchronized_e2e_ms": 1.256737 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.058305, + "submission_ms": 1.058305, + "synchronized_e2e_ms": 1.081281 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.07504, + "submission_ms": 0.07504, + "synchronized_e2e_ms": 0.105216 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.269538, + "submission_ms": 2.269538, + "synchronized_e2e_ms": 2.296898 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.613215, + "evolution_kernel_ms": 0.196895, + "evolution_speedup": 3.1144, + "evolution_tflops": 19.9389, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_68cf_d416_overlap_b2_n2304_k1024_d416", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 9962, + "measurement_schedule_sha256": "5d35856c2357a7c592c3ba1a2ae782f896354ae1d9065aaff71b1c374b6a43fb", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 2844, + "public_pair_count": 2137, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 2844, + "baseline_public_raw": 2137, + "candidate_precomputed": 2844, + "candidate_public_raw": 2137 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 1994 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:adjacent_68cf_d416_overlap_b2_n2304_k1024_d416", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.9258635214827295, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.8029693782864211, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.073281878429081, + "including_init_synchronized_e2e_speedup": 13.187455743121626, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.706865423255624, + "including_init_synchronized_e2e_speedup": 12.909754893482443, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.0356648458470743, + "including_init_synchronized_e2e_speedup": 10.731761818207014, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.8305869234974284, + "including_init_synchronized_e2e_speedup": 4.818478185304285, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.7047163362952835, + "hot_synchronized_e2e_speedup": 1.8029693782864211, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 6841601, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_68cf_d416_overlap_b2_n2304_k1024_d416", + "source": "guard_overlap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 480, + "K": 512, + "N": 1664, + "baseline_07cf_adapter_bench_iters": 2557, + "baseline_07cf_adapter_gpu_span_ms": 0.06256, + "baseline_07cf_adapter_host_enqueue_ms": 0.151136, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.028928, + "baseline_07cf_adapter_kernel_sum_ms": 0.033632, + "baseline_07cf_adapter_submission_ms": 0.151136, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.169728, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.033632 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.204993, + "synchronized_e2e_ms": 0.226689 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.06256 + }, + "host_enqueue_ms": { + "median": 0.151136 + }, + "inter_kernel_gap_ms": { + "median": 0.028928 + }, + "kernel_sum_ms": { + "median": 0.033632 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2557, + "submission_ms": { + "median": 0.151136 + }, + "synchronized_e2e_ms": { + "median": 0.169728 + } + }, + "baseline_07cf_precomputed_bench_iters": 4302, + "baseline_07cf_precomputed_gpu_span_ms": 0.023168, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042848, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.023168, + "baseline_07cf_precomputed_submission_ms": 0.042848, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.071584, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.023168 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.045376, + "synchronized_e2e_ms": 0.066112 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.023168 + }, + "host_enqueue_ms": { + "median": 0.042848 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.023168 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4302, + "submission_ms": { + "median": 0.042848 + }, + "synchronized_e2e_ms": { + "median": 0.071584 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.700276243093923, + "submission": 3.527259148618372, + "synchronized_e2e": 2.3710326329906124 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.669992, + "after_init_synchronized_e2e_ms_per_call": 7.69652, + "including_init_host_enqueue_ms_per_call": 43.22126, + "including_init_synchronized_e2e_ms_per_call": 43.327757, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9030216, + "after_init_synchronized_e2e_ms_per_call": 0.9224072, + "including_init_host_enqueue_ms_per_call": 4.458148400000001, + "including_init_synchronized_e2e_ms_per_call": 4.4855309, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22632455999999998, + "after_init_synchronized_e2e_ms_per_call": 0.24499592, + "including_init_host_enqueue_ms_per_call": 0.58183724, + "including_init_synchronized_e2e_ms_per_call": 0.60130829, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15865485599999998, + "after_init_synchronized_e2e_ms_per_call": 0.177254792, + "including_init_host_enqueue_ms_per_call": 0.19420612399999998, + "including_init_synchronized_e2e_ms_per_call": 0.21288602899999998, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.669992, + "synchronized_e2e_ms": 7.69652, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.331392, + "median": 0.06256, + "min": 0.052512, + "p90": 0.087008, + "sample_count": 2557 + }, + "host_enqueue_ms": { + "max": 52.401687, + "median": 0.151136, + "min": 0.127232, + "p90": 0.2429252, + "sample_count": 2557 + }, + "sample_count": 2557, + "synchronized_e2e_ms": { + "max": 52.529079, + "median": 0.169728, + "min": 0.144128, + "p90": 0.2691526000000001, + "sample_count": 2557 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 4302, + "candidate_precomputed_gpu_span_ms": 0.026304, + "candidate_precomputed_host_enqueue_ms": 0.053040000000000004, + "candidate_precomputed_inter_kernel_gap_ms": 0.002336, + "candidate_precomputed_kernel_sum_ms": 0.023968, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.053040000000000004, + "candidate_precomputed_synchronized_e2e_ms": 0.064448, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.023968 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.04736, + "synchronized_e2e_ms": 0.062336 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.026304 + }, + "host_enqueue_ms": { + "median": 0.053040000000000004 + }, + "inter_kernel_gap_ms": { + "median": 0.002336 + }, + "kernel_sum_ms": { + "median": 0.023968 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4302, + "submission_ms": { + "median": 0.053040000000000004 + }, + "synchronized_e2e_ms": { + "median": 0.064448 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295c3110", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc046640e0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.489051094890511, + "submission": 0.8458521870286576, + "synchronized_e2e": 1.3222442899702085 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2557, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.039168, + "candidate_public_raw_host_enqueue_ms": 0.044864, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.039392, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.044864, + "candidate_public_raw_synchronized_e2e_ms": 0.085216, + "candidate_public_raw_tflops_from_gpu_span": 83.52627450980393, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.039168 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.071552, + "synchronized_e2e_ms": 0.096224 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.039168 + }, + "host_enqueue_ms": { + "median": 0.044864 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.039392 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2557, + "submission_ms": { + "median": 0.044864 + }, + "synchronized_e2e_ms": { + "median": 0.085216 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.300418, + "after_init_synchronized_e2e_ms_per_call": 2.325602, + "including_init_host_enqueue_ms_per_call": 38.138438, + "including_init_synchronized_e2e_ms_per_call": 486.425815, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.27041940000000003, + "after_init_synchronized_e2e_ms_per_call": 0.3092546, + "including_init_host_enqueue_ms_per_call": 3.8542214, + "including_init_synchronized_e2e_ms_per_call": 48.7192759, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06741954, + "after_init_synchronized_e2e_ms_per_call": 0.10761986, + "including_init_host_enqueue_ms_per_call": 0.42579974, + "including_init_synchronized_e2e_ms_per_call": 4.9486219899999995, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.047119554, + "after_init_synchronized_e2e_ms_per_call": 0.08745638600000001, + "including_init_host_enqueue_ms_per_call": 0.08295757399999999, + "including_init_synchronized_e2e_ms_per_call": 0.571556599, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.300418, + "synchronized_e2e_ms": 2.325602, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.039904, + "median": 0.039168, + "min": 0.038752, + "p90": 0.039392, + "sample_count": 2557 + }, + "host_enqueue_ms": { + "max": 0.529664, + "median": 0.044864, + "min": 0.035712, + "p90": 0.07740800000000003, + "sample_count": 2557 + }, + "sample_count": 2557, + "synchronized_e2e_ms": { + "max": 0.638432, + "median": 0.085216, + "min": 0.07728, + "p90": 0.1107776, + "sample_count": 2557 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.067232, + "submission_ms": 0.067232, + "synchronized_e2e_ms": 0.086688 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.045376, + "submission_ms": 0.045376, + "synchronized_e2e_ms": 0.066112 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.204993, + "submission_ms": 0.204993, + "synchronized_e2e_ms": 0.226689 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.669992, + "submission_ms": 7.669992, + "synchronized_e2e_ms": 7.69652 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.074016, + "submission_ms": 0.074016, + "synchronized_e2e_ms": 0.09168 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.04736, + "submission_ms": 0.04736, + "synchronized_e2e_ms": 0.062336 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.342209, + "submission_ms": 1.342209, + "synchronized_e2e_ms": 1.366465 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.122849, + "submission_ms": 1.122849, + "synchronized_e2e_ms": 1.144417 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.071552, + "submission_ms": 0.071552, + "synchronized_e2e_ms": 0.096224 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.300418, + "submission_ms": 2.300418, + "synchronized_e2e_ms": 2.325602 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.539839, + "evolution_kernel_ms": 0.18536, + "evolution_speedup": 2.9124, + "evolution_tflops": 17.6497, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_68cf_d480_boundary_b4_n1664_k512_d480", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 13718, + "measurement_schedule_sha256": "955d45eef4c0d8c96871ee7b852745cd2e53da6af66125f76322e8155bedd63d", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 4302, + "public_pair_count": 2557, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 4302, + "baseline_public_raw": 2557, + "candidate_precomputed": 4302, + "candidate_public_raw": 2557 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2746 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:adjacent_68cf_d480_boundary_b4_n1664_k512_d480", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.8807785888077859, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.9917386406308673, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.3094742780578965, + "including_init_synchronized_e2e_speedup": 0.08907372031642687, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.9826789965290734, + "including_init_synchronized_e2e_speedup": 0.09206891558090664, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.276493576557338, + "including_init_synchronized_e2e_speedup": 0.12151024895720516, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.0267792908799134, + "including_init_synchronized_e2e_speedup": 0.3724671001480292, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.5972222222222223, + "hot_synchronized_e2e_speedup": 1.9917386406308673, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 6848001, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_68cf_d480_boundary_b4_n1664_k512_d480", + "source": "guard_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 48, + "K": 512, + "N": 2304, + "baseline_07cf_adapter_bench_iters": 6028, + "baseline_07cf_adapter_gpu_span_ms": 0.0594725, + "baseline_07cf_adapter_host_enqueue_ms": 0.162113, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.042176, + "baseline_07cf_adapter_kernel_sum_ms": 0.01728, + "baseline_07cf_adapter_submission_ms": 0.162113, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.18219200000000002, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.01728 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.237952, + "synchronized_e2e_ms": 0.253984 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.0594725 + }, + "host_enqueue_ms": { + "median": 0.162113 + }, + "inter_kernel_gap_ms": { + "median": 0.042176 + }, + "kernel_sum_ms": { + "median": 0.01728 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 6028, + "submission_ms": { + "median": 0.162113 + }, + "synchronized_e2e_ms": { + "median": 0.18219200000000002 + } + }, + "baseline_07cf_precomputed_bench_iters": 8597, + "baseline_07cf_precomputed_gpu_span_ms": 0.01168, + "baseline_07cf_precomputed_host_enqueue_ms": 0.047904, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.01168, + "baseline_07cf_precomputed_submission_ms": 0.047904, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.066592, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.01168 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.038944, + "synchronized_e2e_ms": 0.054816 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.01168 + }, + "host_enqueue_ms": { + "median": 0.047904 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.01168 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 8597, + "submission_ms": { + "median": 0.047904 + }, + "synchronized_e2e_ms": { + "median": 0.066592 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 5.091823630136986, + "submission": 3.384122411489646, + "synchronized_e2e": 2.735944257568477 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 562.03047, + "after_init_synchronized_e2e_ms_per_call": 562.115494, + "including_init_host_enqueue_ms_per_call": 598.964908, + "including_init_synchronized_e2e_ms_per_call": 1017.226813, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 56.34894870000001, + "after_init_synchronized_e2e_ms_per_call": 56.3755222, + "including_init_host_enqueue_ms_per_call": 60.042392500000005, + "including_init_synchronized_e2e_ms_per_call": 101.8866541, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 5.78079657, + "after_init_synchronized_e2e_ms_per_call": 5.801525020000001, + "including_init_host_enqueue_ms_per_call": 6.15014095, + "including_init_synchronized_e2e_ms_per_call": 10.35263821, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.723981357, + "after_init_synchronized_e2e_ms_per_call": 0.744125302, + "including_init_host_enqueue_ms_per_call": 0.760915795, + "including_init_synchronized_e2e_ms_per_call": 1.199236621, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 562.03047, + "synchronized_e2e_ms": 562.115494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 74.845485, + "median": 0.0594725, + "min": 0.046752, + "p90": 0.08080960000000001, + "sample_count": 6028 + }, + "host_enqueue_ms": { + "max": 80.606099, + "median": 0.162113, + "min": 0.128256, + "p90": 0.20504030000000004, + "sample_count": 6028 + }, + "sample_count": 6028, + "synchronized_e2e_ms": { + "max": 80.814835, + "median": 0.18219200000000002, + "min": 0.144672, + "p90": 0.2289696, + "sample_count": 6028 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 8597, + "candidate_precomputed_gpu_span_ms": 0.016064, + "candidate_precomputed_host_enqueue_ms": 0.059904, + "candidate_precomputed_inter_kernel_gap_ms": 0.00672, + "candidate_precomputed_kernel_sum_ms": 0.009344, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.059904, + "candidate_precomputed_synchronized_e2e_ms": 0.072384, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.009344 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.04304, + "synchronized_e2e_ms": 0.055328 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.016064 + }, + "host_enqueue_ms": { + "median": 0.059904 + }, + "inter_kernel_gap_ms": { + "median": 0.00672 + }, + "kernel_sum_ms": { + "median": 0.009344 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 8597, + "submission_ms": { + "median": 0.059904 + }, + "synchronized_e2e_ms": { + "median": 0.072384 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295c3230", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04d61d60" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.043886952191235, + "submission": 0.9353632478632479, + "synchronized_e2e": 1.032714412024757 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 6028, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.016769, + "candidate_public_raw_host_enqueue_ms": 0.056032, + "candidate_public_raw_inter_kernel_gap_ms": 0.000128, + "candidate_public_raw_kernel_sum_ms": 0.016641, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.056032, + "candidate_public_raw_synchronized_e2e_ms": 0.074752, + "candidate_public_raw_tflops_from_gpu_span": 27.013228695807744, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.016641 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.054112, + "synchronized_e2e_ms": 0.06848 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.016769 + }, + "host_enqueue_ms": { + "median": 0.056032 + }, + "inter_kernel_gap_ms": { + "median": 0.000128 + }, + "kernel_sum_ms": { + "median": 0.016641 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 6028, + "submission_ms": { + "median": 0.056032 + }, + "synchronized_e2e_ms": { + "median": 0.074752 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.408036, + "after_init_synchronized_e2e_ms_per_call": 3.431748, + "including_init_host_enqueue_ms_per_call": 40.766059, + "including_init_synchronized_e2e_ms_per_call": 40.867371, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.3912324, + "after_init_synchronized_e2e_ms_per_call": 0.4104516, + "including_init_host_enqueue_ms_per_call": 4.1270347, + "including_init_synchronized_e2e_ms_per_call": 4.1540139, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.08955204, + "after_init_synchronized_e2e_ms_per_call": 0.10832196, + "including_init_host_enqueue_ms_per_call": 0.46313226999999996, + "including_init_synchronized_e2e_ms_per_call": 0.48267818999999995, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.059384004000000004, + "after_init_synchronized_e2e_ms_per_call": 0.078108996, + "including_init_host_enqueue_ms_per_call": 0.09674202700000001, + "including_init_synchronized_e2e_ms_per_call": 0.11554461900000002, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.408036, + "synchronized_e2e_ms": 3.431748, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.01744, + "median": 0.016769, + "min": 0.01632, + "p90": 0.017184, + "sample_count": 6028 + }, + "host_enqueue_ms": { + "max": 25.827227, + "median": 0.056032, + "min": 0.03728, + "p90": 0.07383040000000002, + "sample_count": 6028 + }, + "sample_count": 6028, + "synchronized_e2e_ms": { + "max": 86.285977, + "median": 0.074752, + "min": 0.057312, + "p90": 0.0943136, + "sample_count": 6028 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.069184, + "submission_ms": 0.069184, + "synchronized_e2e_ms": 0.085216 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.038944, + "submission_ms": 0.038944, + "synchronized_e2e_ms": 0.054816 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.237952, + "submission_ms": 0.237952, + "synchronized_e2e_ms": 0.253984 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 562.03047, + "submission_ms": 562.03047, + "synchronized_e2e_ms": 562.115494 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.107744, + "submission_ms": 0.107744, + "synchronized_e2e_ms": 0.125248 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.04304, + "submission_ms": 0.04304, + "synchronized_e2e_ms": 0.055328 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.266017, + "submission_ms": 1.266017, + "synchronized_e2e_ms": 1.289697 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.046754, + "submission_ms": 1.046754, + "synchronized_e2e_ms": 1.06877 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.054112, + "submission_ms": 0.054112, + "synchronized_e2e_ms": 0.06848 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.408036, + "submission_ms": 3.408036, + "synchronized_e2e_ms": 3.431748 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.293984, + "evolution_kernel_ms": 0.170608, + "evolution_speedup": 1.7232, + "evolution_tflops": 2.6551, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_68cf_d48_boundary_b4_n2304_k512_d48", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 29250, + "measurement_schedule_sha256": "c1eacfd7be15b99b90e1de12f26f30161dec93bcb7febff829980c8031adac09", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 8597, + "public_pair_count": 6028, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 8597, + "baseline_public_raw": 6028, + "candidate_precomputed": 8597, + "candidate_public_raw": 6028 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 5852 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:adjacent_68cf_d48_boundary_b4_n2304_k512_d48", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.7270916334661355, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.43728595890411, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 163.79859302023343, + "including_init_synchronized_e2e_speedup": 24.89092858456689, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 137.34998767211528, + "including_init_synchronized_e2e_speedup": 24.52727808638291, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 53.5581614291322, + "including_init_synchronized_e2e_speedup": 21.448324006518714, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 9.526755432882533, + "including_init_synchronized_e2e_speedup": 10.378991521881256, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 3.546574035422506, + "hot_synchronized_e2e_speedup": 2.43728595890411, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 6804801, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_68cf_d48_boundary_b4_n2304_k512_d48", + "source": "guard_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 64, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 112, + "K": 4096, + "N": 384, + "baseline_07cf_adapter_bench_iters": 2860, + "baseline_07cf_adapter_gpu_span_ms": 0.089088, + "baseline_07cf_adapter_host_enqueue_ms": 0.14567999999999998, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.038464, + "baseline_07cf_adapter_kernel_sum_ms": 0.050624, + "baseline_07cf_adapter_submission_ms": 0.14567999999999998, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.1900485, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.050624 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.198144, + "synchronized_e2e_ms": 0.235744 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.089088 + }, + "host_enqueue_ms": { + "median": 0.14567999999999998 + }, + "inter_kernel_gap_ms": { + "median": 0.038464 + }, + "kernel_sum_ms": { + "median": 0.050624 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2860, + "submission_ms": { + "median": 0.14567999999999998 + }, + "synchronized_e2e_ms": { + "median": 0.1900485 + } + }, + "baseline_07cf_precomputed_bench_iters": 3302, + "baseline_07cf_precomputed_gpu_span_ms": 0.054784, + "baseline_07cf_precomputed_host_enqueue_ms": 0.04256, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.054784, + "baseline_07cf_precomputed_submission_ms": 0.04256, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.104, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.054784 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.043392, + "synchronized_e2e_ms": 0.0896 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.054784 + }, + "host_enqueue_ms": { + "median": 0.04256 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.054784 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3302, + "submission_ms": { + "median": 0.04256 + }, + "synchronized_e2e_ms": { + "median": 0.104 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.6261682242990654, + "submission": 3.422932330827067, + "synchronized_e2e": 1.8273894230769232 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.568168, + "after_init_synchronized_e2e_ms_per_call": 7.61604, + "including_init_host_enqueue_ms_per_call": 41.796075, + "including_init_synchronized_e2e_ms_per_call": 41.961483, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8879287999999999, + "after_init_synchronized_e2e_ms_per_call": 0.93264765, + "including_init_host_enqueue_ms_per_call": 4.3107195, + "including_init_synchronized_e2e_ms_per_call": 4.3671919500000005, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.21990488, + "after_init_synchronized_e2e_ms_per_call": 0.264308415, + "including_init_host_enqueue_ms_per_call": 0.56218395, + "including_init_synchronized_e2e_ms_per_call": 0.607762845, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15310248799999998, + "after_init_synchronized_e2e_ms_per_call": 0.1974744915, + "including_init_host_enqueue_ms_per_call": 0.18733039499999998, + "including_init_synchronized_e2e_ms_per_call": 0.2318199345, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.568168, + "synchronized_e2e_ms": 7.61604, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 45.283216, + "median": 0.089088, + "min": 0.0808, + "p90": 0.11347839999999997, + "sample_count": 2860 + }, + "host_enqueue_ms": { + "max": 45.716911, + "median": 0.14567999999999998, + "min": 0.121952, + "p90": 0.23626239999999998, + "sample_count": 2860 + }, + "sample_count": 2860, + "synchronized_e2e_ms": { + "max": 45.821103, + "median": 0.1900485, + "min": 0.167872, + "p90": 0.2702760999999999, + "sample_count": 2860 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3302, + "candidate_precomputed_gpu_span_ms": 0.0304, + "candidate_precomputed_host_enqueue_ms": 0.038016, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.0304, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.038016, + "candidate_precomputed_synchronized_e2e_ms": 0.0651845, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.0304 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.032256, + "synchronized_e2e_ms": 0.05984 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.0304 + }, + "host_enqueue_ms": { + "median": 0.038016 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.0304 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3302, + "submission_ms": { + "median": 0.038016 + }, + "synchronized_e2e_ms": { + "median": 0.0651845 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb9e51cfe0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb9e51f470" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.1505263157894736, + "submission": 1.0867003367003367, + "synchronized_e2e": 1.1973398584019206 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 2860, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.034976, + "candidate_public_raw_host_enqueue_ms": 0.041312, + "candidate_public_raw_inter_kernel_gap_ms": 0.000192, + "candidate_public_raw_kernel_sum_ms": 0.034784, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.041312, + "candidate_public_raw_synchronized_e2e_ms": 0.078048, + "candidate_public_raw_tflops_from_gpu_span": 10.073236962488563, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.034784 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.072864, + "synchronized_e2e_ms": 0.09664 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.034976 + }, + "host_enqueue_ms": { + "median": 0.041312 + }, + "inter_kernel_gap_ms": { + "median": 0.000192 + }, + "kernel_sum_ms": { + "median": 0.034784 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2860, + "submission_ms": { + "median": 0.041312 + }, + "synchronized_e2e_ms": { + "median": 0.078048 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.159298, + "after_init_synchronized_e2e_ms_per_call": 2.182658, + "including_init_host_enqueue_ms_per_call": 36.736838, + "including_init_synchronized_e2e_ms_per_call": 450.27045, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2531106, + "after_init_synchronized_e2e_ms_per_call": 0.288509, + "including_init_host_enqueue_ms_per_call": 3.7108646, + "including_init_synchronized_e2e_ms_per_call": 45.097288199999994, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.062491859999999996, + "after_init_synchronized_e2e_ms_per_call": 0.09909410000000002, + "including_init_host_enqueue_ms_per_call": 0.40826726, + "including_init_synchronized_e2e_ms_per_call": 4.57997202, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.043429986, + "after_init_synchronized_e2e_ms_per_call": 0.08015261000000001, + "including_init_host_enqueue_ms_per_call": 0.078007526, + "including_init_synchronized_e2e_ms_per_call": 0.528240402, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.159298, + "synchronized_e2e_ms": 2.182658, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.036096, + "median": 0.034976, + "min": 0.034401, + "p90": 0.035264, + "sample_count": 2860 + }, + "host_enqueue_ms": { + "max": 20.547157, + "median": 0.041312, + "min": 0.031264, + "p90": 0.0695392, + "sample_count": 2860 + }, + "sample_count": 2860, + "synchronized_e2e_ms": { + "max": 25.730714, + "median": 0.078048, + "min": 0.070176, + "p90": 0.1007424, + "sample_count": 2860 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.062304, + "submission_ms": 0.062304, + "synchronized_e2e_ms": 0.10752 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.043392, + "submission_ms": 0.043392, + "synchronized_e2e_ms": 0.0896 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.198144, + "submission_ms": 0.198144, + "synchronized_e2e_ms": 0.235744 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.568168, + "submission_ms": 7.568168, + "synchronized_e2e_ms": 7.61604 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.057377, + "submission_ms": 0.057377, + "synchronized_e2e_ms": 0.077185 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.032256, + "submission_ms": 0.032256, + "synchronized_e2e_ms": 0.05984 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.019201, + "submission_ms": 1.019201, + "synchronized_e2e_ms": 1.041826 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 0.901024, + "submission_ms": 0.901024, + "synchronized_e2e_ms": 0.922528 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.072864, + "submission_ms": 0.072864, + "synchronized_e2e_ms": 0.09664 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.159298, + "submission_ms": 2.159298, + "synchronized_e2e_ms": 2.182658 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 1.266206, + "evolution_kernel_ms": 0.195296, + "evolution_speedup": 6.4835, + "evolution_tflops": 1.804, + "expected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_8f09_d112_small_highk_b1_n384_k4096_d112", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 12324, + "measurement_schedule_sha256": "0746ad5e36b5b0664b83c855357e99a7aed793547ad0d757cee98df2b271e52f", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3302, + "public_pair_count": 2860, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3302, + "baseline_public_raw": 2860, + "candidate_precomputed": 3302, + "candidate_public_raw": 2860 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2466 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:adjacent_8f09_d112_small_highk_b1_n384_k4096_d112", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.8021052631578947, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.4350207564575643, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.4893418941492436, + "including_init_synchronized_e2e_speedup": 0.09319173176920671, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 3.232646641872524, + "including_init_synchronized_e2e_speedup": 0.09683934720491688, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.667246738201366, + "including_init_synchronized_e2e_speedup": 0.132700121823015, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.4637312683891386, + "including_init_synchronized_e2e_speedup": 0.43885309344437456, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.54711802378774, + "hot_synchronized_e2e_speedup": 2.4350207564575643, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 8091121, + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_8f09_d112_small_highk_b1_n384_k4096_d112", + "source": "heldout_neighborhood", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 128, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 4, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 128, + "K": 512, + "N": 4480, + "baseline_07cf_adapter_bench_iters": 4679, + "baseline_07cf_adapter_gpu_span_ms": 0.062112, + "baseline_07cf_adapter_host_enqueue_ms": 0.16096, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.035648, + "baseline_07cf_adapter_kernel_sum_ms": 0.026464, + "baseline_07cf_adapter_submission_ms": 0.16096, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.181504, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.026464 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.243776, + "synchronized_e2e_ms": 0.265472 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.062112 + }, + "host_enqueue_ms": { + "median": 0.16096 + }, + "inter_kernel_gap_ms": { + "median": 0.035648 + }, + "kernel_sum_ms": { + "median": 0.026464 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 4679, + "submission_ms": { + "median": 0.16096 + }, + "synchronized_e2e_ms": { + "median": 0.181504 + } + }, + "baseline_07cf_precomputed_bench_iters": 12808, + "baseline_07cf_precomputed_gpu_span_ms": 0.014336, + "baseline_07cf_precomputed_host_enqueue_ms": 0.04496, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.014336, + "baseline_07cf_precomputed_submission_ms": 0.04496, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.0646245, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.014336 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.045249, + "synchronized_e2e_ms": 0.062945 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.014336 + }, + "host_enqueue_ms": { + "median": 0.04496 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.014336 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 12808, + "submission_ms": { + "median": 0.04496 + }, + "synchronized_e2e_ms": { + "median": 0.0646245 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 4.332589285714286, + "submission": 3.580071174377224, + "synchronized_e2e": 2.808594263785406 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 6.601031, + "after_init_synchronized_e2e_ms_per_call": 6.628775, + "including_init_host_enqueue_ms_per_call": 41.050475, + "including_init_synchronized_e2e_ms_per_call": 490.670269, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8049671, + "after_init_synchronized_e2e_ms_per_call": 0.8262311, + "including_init_host_enqueue_ms_per_call": 4.2499115, + "including_init_synchronized_e2e_ms_per_call": 49.2303805, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22536071, + "after_init_synchronized_e2e_ms_per_call": 0.24597671000000002, + "including_init_host_enqueue_ms_per_call": 0.56985515, + "including_init_synchronized_e2e_ms_per_call": 5.08639165, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.16740007099999998, + "after_init_synchronized_e2e_ms_per_call": 0.187951271, + "including_init_host_enqueue_ms_per_call": 0.201849515, + "including_init_synchronized_e2e_ms_per_call": 0.6719927649999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 6.601031, + "synchronized_e2e_ms": 6.628775, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.496864, + "median": 0.062112, + "min": 0.050688, + "p90": 0.08336, + "sample_count": 4679 + }, + "host_enqueue_ms": { + "max": 5.480998, + "median": 0.16096, + "min": 0.131616, + "p90": 0.20878079999999993, + "sample_count": 4679 + }, + "sample_count": 4679, + "synchronized_e2e_ms": { + "max": 15.713232, + "median": 0.181504, + "min": 0.14864, + "p90": 0.23389439999999997, + "sample_count": 4679 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 12808, + "candidate_precomputed_gpu_span_ms": 0.008032, + "candidate_precomputed_host_enqueue_ms": 0.041984, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.008032, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.041984, + "candidate_precomputed_synchronized_e2e_ms": 0.05408, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.008032 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.036032, + "synchronized_e2e_ms": 0.052 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.008032 + }, + "host_enqueue_ms": { + "median": 0.041984 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.008032 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 12808, + "submission_ms": { + "median": 0.041984 + }, + "synchronized_e2e_ms": { + "median": 0.05408 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01993890", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01990050" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.6812749003984067, + "submission": 1.2195360137195124, + "synchronized_e2e": 1.3349112426035503 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 4679, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.021536, + "candidate_public_raw_host_enqueue_ms": 0.051201, + "candidate_public_raw_inter_kernel_gap_ms": 0.00016, + "candidate_public_raw_kernel_sum_ms": 0.021376, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.051201, + "candidate_public_raw_synchronized_e2e_ms": 0.072192, + "candidate_public_raw_tflops_from_gpu_span": 109.06436849925707, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.021376 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.261633, + "synchronized_e2e_ms": 0.286785 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.021536 + }, + "host_enqueue_ms": { + "median": 0.051201 + }, + "inter_kernel_gap_ms": { + "median": 0.00016 + }, + "kernel_sum_ms": { + "median": 0.021376 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 4679, + "submission_ms": { + "median": 0.051201 + }, + "synchronized_e2e_ms": { + "median": 0.072192 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.196226, + "after_init_synchronized_e2e_ms_per_call": 2.223362, + "including_init_host_enqueue_ms_per_call": 37.07031, + "including_init_synchronized_e2e_ms_per_call": 37.166438, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2657035, + "after_init_synchronized_e2e_ms_per_call": 0.287309, + "including_init_host_enqueue_ms_per_call": 3.7531118999999995, + "including_init_synchronized_e2e_ms_per_call": 3.7816166000000004, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07265125, + "after_init_synchronized_e2e_ms_per_call": 0.09370370000000001, + "including_init_host_enqueue_ms_per_call": 0.42139209, + "including_init_synchronized_e2e_ms_per_call": 0.44313446, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.053346025000000005, + "after_init_synchronized_e2e_ms_per_call": 0.07434317, + "including_init_host_enqueue_ms_per_call": 0.088220109, + "including_init_synchronized_e2e_ms_per_call": 0.109286246, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.196226, + "synchronized_e2e_ms": 2.223362, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.022336, + "median": 0.021536, + "min": 0.02096, + "p90": 0.021792, + "sample_count": 4679 + }, + "host_enqueue_ms": { + "max": 25.543514, + "median": 0.051201, + "min": 0.038592, + "p90": 0.07367679999999996, + "sample_count": 4679 + }, + "sample_count": 4679, + "synchronized_e2e_ms": { + "max": 31.865153, + "median": 0.072192, + "min": 0.060576, + "p90": 0.0959816, + "sample_count": 4679 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.067008, + "submission_ms": 0.067008, + "synchronized_e2e_ms": 0.084768 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.045249, + "submission_ms": 0.045249, + "synchronized_e2e_ms": 0.062945 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.243776, + "submission_ms": 0.243776, + "synchronized_e2e_ms": 0.265472 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 6.601031, + "submission_ms": 6.601031, + "synchronized_e2e_ms": 6.628775 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.055648, + "submission_ms": 0.055648, + "synchronized_e2e_ms": 0.071264 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.036032, + "submission_ms": 0.036032, + "synchronized_e2e_ms": 0.052 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.196417, + "submission_ms": 1.196417, + "synchronized_e2e_ms": 1.217985 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.094593, + "submission_ms": 1.094593, + "synchronized_e2e_ms": 1.115681 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.261633, + "submission_ms": 0.261633, + "synchronized_e2e_ms": 0.286785 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.196226, + "submission_ms": 2.196226, + "synchronized_e2e_ms": 2.223362 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.325535, + "evolution_kernel_ms": 0.145856, + "evolution_speedup": 2.2319, + "evolution_tflops": 16.1036, + "expected_route": "aligned_weave_v10_fallback", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_8f09_d128_fallback_neighbor_b4_n4480_k512_d128", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 34974, + "measurement_schedule_sha256": "2e5d8b08d5ec88470d4becc5dd6838f766fe2a09962d16731e4e443bb21530f1", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 12808, + "public_pair_count": 4679, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 12808, + "baseline_public_raw": 4679, + "candidate_precomputed": 12808, + "candidate_public_raw": 4679 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 6996 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:adjacent_8f09_d128_fallback_neighbor_b4_n4480_k512_d128", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.7848605577689245, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.5141843971631204, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.9814195798974708, + "including_init_synchronized_e2e_speedup": 13.201971870427831, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.875757807795788, + "including_init_synchronized_e2e_speedup": 13.01834260511761, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.62504799703747, + "including_init_synchronized_e2e_speedup": 11.47821284311764, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.528157879197242, + "including_init_synchronized_e2e_speedup": 6.148923488505589, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.8841010401188707, + "hot_synchronized_e2e_speedup": 2.5141843971631204, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 8091281, + "selected_route": "aligned_weave_v10_fallback", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_8f09_d128_fallback_neighbor_b4_n4480_k512_d128", + "source": "guard_miss_fallback", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 224, + "K": 256, + "N": 1536, + "baseline_07cf_adapter_bench_iters": 5037, + "baseline_07cf_adapter_gpu_span_ms": 0.053408, + "baseline_07cf_adapter_host_enqueue_ms": 0.151104, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.0344, + "baseline_07cf_adapter_kernel_sum_ms": 0.018976, + "baseline_07cf_adapter_submission_ms": 0.151104, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.170496, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.018976 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.209376, + "synchronized_e2e_ms": 0.23216 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.053408 + }, + "host_enqueue_ms": { + "median": 0.151104 + }, + "inter_kernel_gap_ms": { + "median": 0.0344 + }, + "kernel_sum_ms": { + "median": 0.018976 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 5037, + "submission_ms": { + "median": 0.151104 + }, + "synchronized_e2e_ms": { + "median": 0.170496 + } + }, + "baseline_07cf_precomputed_bench_iters": 9705, + "baseline_07cf_precomputed_gpu_span_ms": 0.010368, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042048, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.010368, + "baseline_07cf_precomputed_submission_ms": 0.042048, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.057856, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.010368 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.044032, + "synchronized_e2e_ms": 0.060832 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.010368 + }, + "host_enqueue_ms": { + "median": 0.042048 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.010368 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 9705, + "submission_ms": { + "median": 0.042048 + }, + "synchronized_e2e_ms": { + "median": 0.057856 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 5.151234567901234, + "submission": 3.5936073059360725, + "synchronized_e2e": 2.946902654867257 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.298024, + "after_init_synchronized_e2e_ms_per_call": 7.32468, + "including_init_host_enqueue_ms_per_call": 42.849292, + "including_init_synchronized_e2e_ms_per_call": 42.955917, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8657959999999999, + "after_init_synchronized_e2e_ms_per_call": 0.8859144000000001, + "including_init_host_enqueue_ms_per_call": 4.4209228, + "including_init_synchronized_e2e_ms_per_call": 4.4490381, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2225732, + "after_init_synchronized_e2e_ms_per_call": 0.24203784000000003, + "including_init_host_enqueue_ms_per_call": 0.57808588, + "including_init_synchronized_e2e_ms_per_call": 0.59835021, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15825092, + "after_init_synchronized_e2e_ms_per_call": 0.177650184, + "including_init_host_enqueue_ms_per_call": 0.193802188, + "including_init_synchronized_e2e_ms_per_call": 0.213281421, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.298024, + "synchronized_e2e_ms": 7.32468, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.378304, + "median": 0.053408, + "min": 0.043712, + "p90": 0.08020480000000002, + "sample_count": 5037 + }, + "host_enqueue_ms": { + "max": 40.191977, + "median": 0.151104, + "min": 0.125121, + "p90": 0.2470406000000001, + "sample_count": 5037 + }, + "sample_count": 5037, + "synchronized_e2e_ms": { + "max": 40.273897, + "median": 0.170496, + "min": 0.140353, + "p90": 0.27318460000000017, + "sample_count": 5037 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 9705, + "candidate_precomputed_gpu_span_ms": 0.010528, + "candidate_precomputed_host_enqueue_ms": 0.039744, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.010528, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.039744, + "candidate_precomputed_synchronized_e2e_ms": 0.05104, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.010528 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.060928, + "synchronized_e2e_ms": 0.08112 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.010528 + }, + "host_enqueue_ms": { + "median": 0.039744 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.010528 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 9705, + "submission_ms": { + "median": 0.039744 + }, + "synchronized_e2e_ms": { + "median": 0.05104 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc029a16a0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb937a1760" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.905775075987842, + "submission": 1.1344605475040257, + "synchronized_e2e": 1.2771159874608151 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 5037, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.020064, + "candidate_public_raw_host_enqueue_ms": 0.045088, + "candidate_public_raw_inter_kernel_gap_ms": 0.00016, + "candidate_public_raw_kernel_sum_ms": 0.019904, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.045088, + "candidate_public_raw_synchronized_e2e_ms": 0.065184, + "candidate_public_raw_tflops_from_gpu_span": 35.11977033492823, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.019904 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.223297, + "synchronized_e2e_ms": 0.246849 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.020064 + }, + "host_enqueue_ms": { + "median": 0.045088 + }, + "inter_kernel_gap_ms": { + "median": 0.00016 + }, + "kernel_sum_ms": { + "median": 0.019904 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 5037, + "submission_ms": { + "median": 0.045088 + }, + "synchronized_e2e_ms": { + "median": 0.065184 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 37.841671, + "after_init_synchronized_e2e_ms_per_call": 37.869063, + "including_init_host_enqueue_ms_per_call": 73.67969099999999, + "including_init_synchronized_e2e_ms_per_call": 521.969276, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 3.8247462999999997, + "after_init_synchronized_e2e_ms_per_call": 3.8455718999999995, + "including_init_host_enqueue_ms_per_call": 7.4085483, + "including_init_synchronized_e2e_ms_per_call": 52.2555932, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.42305382999999996, + "after_init_synchronized_e2e_ms_per_call": 0.4432227899999999, + "including_init_host_enqueue_ms_per_call": 0.7814340299999999, + "including_init_synchronized_e2e_ms_per_call": 5.284224920000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.082884583, + "after_init_synchronized_e2e_ms_per_call": 0.102987879, + "including_init_host_enqueue_ms_per_call": 0.118722603, + "including_init_synchronized_e2e_ms_per_call": 0.5870880920000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 37.841671, + "synchronized_e2e_ms": 37.869063, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.020896, + "median": 0.020064, + "min": 0.019488, + "p90": 0.020416, + "sample_count": 5037 + }, + "host_enqueue_ms": { + "max": 43.136173, + "median": 0.045088, + "min": 0.034112, + "p90": 0.077408, + "sample_count": 5037 + }, + "sample_count": 5037, + "synchronized_e2e_ms": { + "max": 43.262445, + "median": 0.065184, + "min": 0.056224, + "p90": 0.0991044, + "sample_count": 5037 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.070752, + "submission_ms": 0.070752, + "synchronized_e2e_ms": 0.090048 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.044032, + "submission_ms": 0.044032, + "synchronized_e2e_ms": 0.060832 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.209376, + "submission_ms": 0.209376, + "synchronized_e2e_ms": 0.23216 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.298024, + "submission_ms": 7.298024, + "synchronized_e2e_ms": 7.32468 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.055616, + "submission_ms": 0.055616, + "synchronized_e2e_ms": 0.073088 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.060928, + "submission_ms": 0.060928, + "synchronized_e2e_ms": 0.08112 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.168737, + "submission_ms": 1.168737, + "synchronized_e2e_ms": 1.190881 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.051937, + "submission_ms": 1.051937, + "synchronized_e2e_ms": 1.073345 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.223297, + "submission_ms": 0.223297, + "synchronized_e2e_ms": 0.246849 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 37.841671, + "submission_ms": 37.841671, + "synchronized_e2e_ms": 37.869063 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.317296, + "evolution_kernel_ms": 0.173664, + "evolution_speedup": 1.8271, + "evolution_tflops": 4.0575, + "expected_route": "d224_tmem_abi_repair_d17c_v4", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_8f09_d224_low_n_boundary_b4_n1536_k256_d224", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 29484, + "measurement_schedule_sha256": "f7e3976af5c7b468bcac7b90a7af0104a20ef0645970e759747c817d5e0f0b4e", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 9705, + "public_pair_count": 5037, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 9705, + "baseline_public_raw": 5037, + "candidate_precomputed": 9705, + "candidate_public_raw": 5037 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 5898 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:adjacent_8f09_d224_low_n_boundary_b4_n1536_k256_d224", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.9848024316109424, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.6156111929307806, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.1934212103425955, + "including_init_synchronized_e2e_speedup": 0.0822958725256465, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.23037260075673016, + "including_init_synchronized_e2e_speedup": 0.08513994057960479, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.5460861793681685, + "including_init_synchronized_e2e_speedup": 0.11323329704141358, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.7249620608265948, + "including_init_synchronized_e2e_speedup": 0.3632869136783649, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.6618819776714515, + "hot_synchronized_e2e_speedup": 2.6156111929307806, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 8092241, + "selected_route": "d224_tmem_abi_repair_d17c_v4", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_8f09_d224_low_n_boundary_b4_n1536_k256_d224", + "source": "guard_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 288, + "K": 768, + "N": 2560, + "baseline_07cf_adapter_bench_iters": 3124, + "baseline_07cf_adapter_gpu_span_ms": 0.074304, + "baseline_07cf_adapter_host_enqueue_ms": 0.16561599999999999, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.045632, + "baseline_07cf_adapter_kernel_sum_ms": 0.028672, + "baseline_07cf_adapter_submission_ms": 0.16561599999999999, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.186976, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.028672 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.232416, + "synchronized_e2e_ms": 0.253344 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.074304 + }, + "host_enqueue_ms": { + "median": 0.16561599999999999 + }, + "inter_kernel_gap_ms": { + "median": 0.045632 + }, + "kernel_sum_ms": { + "median": 0.028672 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3124, + "submission_ms": { + "median": 0.16561599999999999 + }, + "synchronized_e2e_ms": { + "median": 0.186976 + } + }, + "baseline_07cf_precomputed_bench_iters": 3939, + "baseline_07cf_precomputed_gpu_span_ms": 0.02544, + "baseline_07cf_precomputed_host_enqueue_ms": 0.051393, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.02544, + "baseline_07cf_precomputed_submission_ms": 0.051393, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.083552, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.02544 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.041696, + "synchronized_e2e_ms": 0.068288 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.02544 + }, + "host_enqueue_ms": { + "median": 0.051393 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.02544 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3939, + "submission_ms": { + "median": 0.051393 + }, + "synchronized_e2e_ms": { + "median": 0.083552 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.920754716981132, + "submission": 3.2225400346350668, + "synchronized_e2e": 2.237839908081195 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.717768, + "after_init_synchronized_e2e_ms_per_call": 7.744776, + "including_init_host_enqueue_ms_per_call": 44.652206, + "including_init_synchronized_e2e_ms_per_call": 462.85609500000004, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9208312, + "after_init_synchronized_e2e_ms_per_call": 0.9427559999999999, + "including_init_host_enqueue_ms_per_call": 4.614275, + "including_init_synchronized_e2e_ms_per_call": 46.453887900000005, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.24113752, + "after_init_synchronized_e2e_ms_per_call": 0.262554, + "including_init_host_enqueue_ms_per_call": 0.6104819, + "including_init_synchronized_e2e_ms_per_call": 4.81366719, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.173168152, + "after_init_synchronized_e2e_ms_per_call": 0.1945338, + "including_init_host_enqueue_ms_per_call": 0.21010259, + "including_init_synchronized_e2e_ms_per_call": 0.649645119, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.717768, + "synchronized_e2e_ms": 7.744776, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.293957, + "median": 0.074304, + "min": 0.063744, + "p90": 0.09378560000000004, + "sample_count": 3124 + }, + "host_enqueue_ms": { + "max": 63.22173, + "median": 0.16561599999999999, + "min": 0.142048, + "p90": 0.20028799999999997, + "sample_count": 3124 + }, + "sample_count": 3124, + "synchronized_e2e_ms": { + "max": 63.348258, + "median": 0.186976, + "min": 0.16288, + "p90": 0.2243616000000001, + "sample_count": 3124 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3939, + "candidate_precomputed_gpu_span_ms": 0.028544, + "candidate_precomputed_host_enqueue_ms": 0.046976, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.028544, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.046976, + "candidate_precomputed_synchronized_e2e_ms": 0.069473, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.028544 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.031904, + "synchronized_e2e_ms": 0.051552 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.028544 + }, + "host_enqueue_ms": { + "median": 0.046976 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.028544 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3939, + "submission_ms": { + "median": 0.046976 + }, + "synchronized_e2e_ms": { + "median": 0.069473 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb6ac30080", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb6ac32510" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.1244394618834082, + "submission": 1.2363866655313351, + "synchronized_e2e": 1.3159644754077127 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 3124, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.032096, + "candidate_public_raw_host_enqueue_ms": 0.0580805, + "candidate_public_raw_inter_kernel_gap_ms": 0.000192, + "candidate_public_raw_kernel_sum_ms": 0.031873, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.0580805, + "candidate_public_raw_synchronized_e2e_ms": 0.091424, + "candidate_public_raw_tflops_from_gpu_span": 35.28358923230309, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.031873 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.275456, + "synchronized_e2e_ms": 0.298048 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.032096 + }, + "host_enqueue_ms": { + "median": 0.0580805 + }, + "inter_kernel_gap_ms": { + "median": 0.000192 + }, + "kernel_sum_ms": { + "median": 0.031873 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3124, + "submission_ms": { + "median": 0.0580805 + }, + "synchronized_e2e_ms": { + "median": 0.091424 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.3313, + "after_init_synchronized_e2e_ms_per_call": 3.354596, + "including_init_host_enqueue_ms_per_call": 40.689322999999995, + "including_init_synchronized_e2e_ms_per_call": 40.790219, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.38540245000000006, + "after_init_synchronized_e2e_ms_per_call": 0.41774120000000003, + "including_init_host_enqueue_ms_per_call": 4.1212047499999995, + "including_init_synchronized_e2e_ms_per_call": 4.161303500000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.090812695, + "after_init_synchronized_e2e_ms_per_call": 0.12405572, + "including_init_host_enqueue_ms_per_call": 0.46439292499999996, + "including_init_synchronized_e2e_ms_per_call": 0.49841195, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.061353719499999994, + "after_init_synchronized_e2e_ms_per_call": 0.094687172, + "including_init_host_enqueue_ms_per_call": 0.09871174249999999, + "including_init_synchronized_e2e_ms_per_call": 0.132122795, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.3313, + "synchronized_e2e_ms": 3.354596, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.032736, + "median": 0.032096, + "min": 0.03152, + "p90": 0.03232, + "sample_count": 3124 + }, + "host_enqueue_ms": { + "max": 15.85352, + "median": 0.0580805, + "min": 0.046656, + "p90": 0.07306910000000003, + "sample_count": 3124 + }, + "sample_count": 3124, + "synchronized_e2e_ms": { + "max": 26.099451, + "median": 0.091424, + "min": 0.080096, + "p90": 0.10425920000000004, + "sample_count": 3124 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.07184, + "submission_ms": 0.07184, + "synchronized_e2e_ms": 0.097664 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.041696, + "submission_ms": 0.041696, + "synchronized_e2e_ms": 0.068288 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.232416, + "submission_ms": 0.232416, + "synchronized_e2e_ms": 0.253344 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.717768, + "submission_ms": 7.717768, + "synchronized_e2e_ms": 7.744776 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.071968, + "submission_ms": 0.071968, + "synchronized_e2e_ms": 0.086816 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.031904, + "submission_ms": 0.031904, + "synchronized_e2e_ms": 0.051552 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.27997, + "submission_ms": 1.27997, + "synchronized_e2e_ms": 1.302786 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.040417, + "submission_ms": 1.040417, + "synchronized_e2e_ms": 1.062017 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.275456, + "submission_ms": 0.275456, + "synchronized_e2e_ms": 0.298048 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.3313, + "submission_ms": 3.3313, + "synchronized_e2e_ms": 3.354596 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.480287, + "evolution_kernel_ms": 0.18432, + "evolution_speedup": 2.6057, + "evolution_tflops": 6.144, + "expected_route": "d288_parent_splitk_hybrid_20260629_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_8f09_d288_k768_overlap_b1_n2560_k768_d288", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 14126, + "measurement_schedule_sha256": "d857d1a850d3bd378f38eb99338d8569ed404853efbdb5c8995fc3b93d459faf", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3939, + "public_pair_count": 3124, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3939, + "baseline_public_raw": 3124, + "candidate_precomputed": 3939, + "candidate_public_raw": 3124 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2826 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:adjacent_8f09_d288_k768_overlap_b1_n2560_k768_d288", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.891255605381166, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.0451522576128807, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.308706026001343, + "including_init_synchronized_e2e_speedup": 11.347232408828205, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.256794398062724, + "including_init_synchronized_e2e_speedup": 11.163302051869083, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.116419944199268, + "including_init_synchronized_e2e_speedup": 9.658009182966019, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.0544894930434716, + "including_init_synchronized_e2e_speedup": 4.916979836825281, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.315054835493519, + "hot_synchronized_e2e_speedup": 2.0451522576128807, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 8092881, + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_8f09_d288_k768_overlap_b1_n2560_k768_d288", + "source": "guard_overlap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 352, + "K": 256, + "N": 4096, + "baseline_07cf_adapter_bench_iters": 2345, + "baseline_07cf_adapter_gpu_span_ms": 0.065088, + "baseline_07cf_adapter_host_enqueue_ms": 0.153728, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.017536, + "baseline_07cf_adapter_kernel_sum_ms": 0.04752, + "baseline_07cf_adapter_submission_ms": 0.153728, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.173696, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.04752 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.201185, + "synchronized_e2e_ms": 0.220257 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.065088 + }, + "host_enqueue_ms": { + "median": 0.153728 + }, + "inter_kernel_gap_ms": { + "median": 0.017536 + }, + "kernel_sum_ms": { + "median": 0.04752 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2345, + "submission_ms": { + "median": 0.153728 + }, + "synchronized_e2e_ms": { + "median": 0.173696 + } + }, + "baseline_07cf_precomputed_bench_iters": 5162, + "baseline_07cf_precomputed_gpu_span_ms": 0.0248, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042656, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.0248, + "baseline_07cf_precomputed_submission_ms": 0.042656, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.074176, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.0248 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.042913, + "synchronized_e2e_ms": 0.067585 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.0248 + }, + "host_enqueue_ms": { + "median": 0.042656 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.0248 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5162, + "submission_ms": { + "median": 0.042656 + }, + "synchronized_e2e_ms": { + "median": 0.074176 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.624516129032258, + "submission": 3.603900975243811, + "synchronized_e2e": 2.3416738567730797 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.15924, + "after_init_synchronized_e2e_ms_per_call": 8.185576, + "including_init_host_enqueue_ms_per_call": 42.387147, + "including_init_synchronized_e2e_ms_per_call": 42.531019, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9542792, + "after_init_synchronized_e2e_ms_per_call": 0.974884, + "including_init_host_enqueue_ms_per_call": 4.3770699, + "including_init_synchronized_e2e_ms_per_call": 4.4094283, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23378312, + "after_init_synchronized_e2e_ms_per_call": 0.25381479999999995, + "including_init_host_enqueue_ms_per_call": 0.5760621899999999, + "including_init_synchronized_e2e_ms_per_call": 0.59726923, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.16173351200000002, + "after_init_synchronized_e2e_ms_per_call": 0.18170788, + "including_init_host_enqueue_ms_per_call": 0.195961419, + "including_init_synchronized_e2e_ms_per_call": 0.21605332299999996, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.15924, + "synchronized_e2e_ms": 8.185576, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.371393, + "median": 0.065088, + "min": 0.055616, + "p90": 0.09159679999999998, + "sample_count": 2345 + }, + "host_enqueue_ms": { + "max": 53.319479, + "median": 0.153728, + "min": 0.126496, + "p90": 0.23690879999999997, + "sample_count": 2345 + }, + "sample_count": 2345, + "synchronized_e2e_ms": { + "max": 53.432023, + "median": 0.173696, + "min": 0.145728, + "p90": 0.2631808, + "sample_count": 2345 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 5162, + "candidate_precomputed_gpu_span_ms": 0.018944, + "candidate_precomputed_host_enqueue_ms": 0.053152, + "candidate_precomputed_inter_kernel_gap_ms": 0.001984, + "candidate_precomputed_kernel_sum_ms": 0.017056, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.053152, + "candidate_precomputed_synchronized_e2e_ms": 0.065248, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.017056 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.043808, + "synchronized_e2e_ms": 0.057856 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.018944 + }, + "host_enqueue_ms": { + "median": 0.053152 + }, + "inter_kernel_gap_ms": { + "median": 0.001984 + }, + "kernel_sum_ms": { + "median": 0.017056 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5162, + "submission_ms": { + "median": 0.053152 + }, + "synchronized_e2e_ms": { + "median": 0.065248 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01dbb920", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01dbade0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.277027027027027, + "submission": 0.8603251053582179, + "synchronized_e2e": 1.3898970083374205 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2345, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.043136, + "candidate_public_raw_host_enqueue_ms": 0.045728, + "candidate_public_raw_inter_kernel_gap_ms": 9.6e-05, + "candidate_public_raw_kernel_sum_ms": 0.04304, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.045728, + "candidate_public_raw_synchronized_e2e_ms": 0.090688, + "candidate_public_raw_tflops_from_gpu_span": 68.45303264094956, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.04304 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.073184, + "synchronized_e2e_ms": 0.100768 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.043136 + }, + "host_enqueue_ms": { + "median": 0.045728 + }, + "inter_kernel_gap_ms": { + "median": 9.6e-05 + }, + "kernel_sum_ms": { + "median": 0.04304 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2345, + "submission_ms": { + "median": 0.045728 + }, + "synchronized_e2e_ms": { + "median": 0.090688 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.720131, + "after_init_synchronized_e2e_ms_per_call": 2.744227, + "including_init_host_enqueue_ms_per_call": 37.297671, + "including_init_synchronized_e2e_ms_per_call": 450.832019, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.31316829999999996, + "after_init_synchronized_e2e_ms_per_call": 0.3560419, + "including_init_host_enqueue_ms_per_call": 3.7709223, + "including_init_synchronized_e2e_ms_per_call": 45.1648211, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07247202999999999, + "after_init_synchronized_e2e_ms_per_call": 0.11722339000000001, + "including_init_host_enqueue_ms_per_call": 0.41824743, + "including_init_synchronized_e2e_ms_per_call": 4.5981013100000006, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.048402403, + "after_init_synchronized_e2e_ms_per_call": 0.093341539, + "including_init_host_enqueue_ms_per_call": 0.08297994299999999, + "including_init_synchronized_e2e_ms_per_call": 0.541429331, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.720131, + "synchronized_e2e_ms": 2.744227, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.044032, + "median": 0.043136, + "min": 0.042176, + "p90": 0.043456, + "sample_count": 2345 + }, + "host_enqueue_ms": { + "max": 38.951881, + "median": 0.045728, + "min": 0.036992, + "p90": 0.07238399999999998, + "sample_count": 2345 + }, + "sample_count": 2345, + "synchronized_e2e_ms": { + "max": 39.055753, + "median": 0.090688, + "min": 0.083264, + "p90": 0.11203260000000001, + "sample_count": 2345 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.070528, + "submission_ms": 0.070528, + "synchronized_e2e_ms": 0.0944 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.042913, + "submission_ms": 0.042913, + "synchronized_e2e_ms": 0.067585 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.201185, + "submission_ms": 0.201185, + "synchronized_e2e_ms": 0.220257 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.15924, + "submission_ms": 8.15924, + "synchronized_e2e_ms": 8.185576 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.067072, + "submission_ms": 0.067072, + "synchronized_e2e_ms": 0.083456 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.043808, + "submission_ms": 0.043808, + "synchronized_e2e_ms": 0.057856 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.387681, + "submission_ms": 1.387681, + "synchronized_e2e_ms": 1.411009 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.163585, + "submission_ms": 1.163585, + "synchronized_e2e_ms": 1.186817 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.073184, + "submission_ms": 0.073184, + "synchronized_e2e_ms": 0.100768 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.720131, + "submission_ms": 2.720131, + "synchronized_e2e_ms": 2.744227 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.306655, + "evolution_kernel_ms": 0.176128, + "evolution_speedup": 1.7411, + "evolution_tflops": 16.765, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_8f09_d352_smallk_large_n_b4_n4096_k256_d352", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 15014, + "measurement_schedule_sha256": "5f369513f8fd6462f9583742f5819792472e7371677fddac5d47c9fce1587d30", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 5162, + "public_pair_count": 2345, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 5162, + "baseline_public_raw": 2345, + "candidate_precomputed": 5162, + "candidate_public_raw": 2345 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 3004 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:adjacent_8f09_d352_smallk_large_n_b4_n4096_k256_d352", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.3091216216216217, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.9153140437544105, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.982834874811741, + "including_init_synchronized_e2e_speedup": 0.09433894933713659, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.7381159352312183, + "including_init_synchronized_e2e_speedup": 0.09762970809154828, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.16522316919857, + "including_init_synchronized_e2e_speedup": 0.12989475214498916, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.9466989932531538, + "including_init_synchronized_e2e_speedup": 0.3990425169633079, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.508902077151335, + "hot_synchronized_e2e_speedup": 1.9153140437544105, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 8093521, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_8f09_d352_smallk_large_n_b4_n4096_k256_d352", + "source": "random_legal", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 3, + "D": 416, + "K": 768, + "N": 3456, + "baseline_07cf_adapter_bench_iters": 2024, + "baseline_07cf_adapter_gpu_span_ms": 0.0900965, + "baseline_07cf_adapter_host_enqueue_ms": 0.1508005, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.024672, + "baseline_07cf_adapter_kernel_sum_ms": 0.065376, + "baseline_07cf_adapter_submission_ms": 0.1508005, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.191296, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.065376 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.178272, + "synchronized_e2e_ms": 0.216384 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.0900965 + }, + "host_enqueue_ms": { + "median": 0.1508005 + }, + "inter_kernel_gap_ms": { + "median": 0.024672 + }, + "kernel_sum_ms": { + "median": 0.065376 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2024, + "submission_ms": { + "median": 0.1508005 + }, + "synchronized_e2e_ms": { + "median": 0.191296 + } + }, + "baseline_07cf_precomputed_bench_iters": 2893, + "baseline_07cf_precomputed_gpu_span_ms": 0.051776, + "baseline_07cf_precomputed_host_enqueue_ms": 0.044544, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.051776, + "baseline_07cf_precomputed_submission_ms": 0.044544, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.10208, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.051776 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.046752, + "synchronized_e2e_ms": 0.094657 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.051776 + }, + "host_enqueue_ms": { + "median": 0.044544 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.051776 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2893, + "submission_ms": { + "median": 0.044544 + }, + "synchronized_e2e_ms": { + "median": 0.10208 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.7401209054388131, + "submission": 3.3854278915229887, + "synchronized_e2e": 1.8739811912225703 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 6.782055, + "after_init_synchronized_e2e_ms_per_call": 6.808455, + "including_init_host_enqueue_ms_per_call": 41.231499, + "including_init_synchronized_e2e_ms_per_call": 490.849949, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.81392595, + "after_init_synchronized_e2e_ms_per_call": 0.8530119, + "including_init_host_enqueue_ms_per_call": 4.2588703500000005, + "including_init_synchronized_e2e_ms_per_call": 49.25716129999999, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.217113045, + "after_init_synchronized_e2e_ms_per_call": 0.25746758999999997, + "including_init_host_enqueue_ms_per_call": 0.561607485, + "including_init_synchronized_e2e_ms_per_call": 5.09788253, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15743175450000002, + "after_init_synchronized_e2e_ms_per_call": 0.197913159, + "including_init_host_enqueue_ms_per_call": 0.1918811985, + "including_init_synchronized_e2e_ms_per_call": 0.681954653, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 6.782055, + "synchronized_e2e_ms": 6.808455, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.298432, + "median": 0.0900965, + "min": 0.081408, + "p90": 0.10993920000000004, + "sample_count": 2024 + }, + "host_enqueue_ms": { + "max": 40.563914, + "median": 0.1508005, + "min": 0.125376, + "p90": 0.1889728, + "sample_count": 2024 + }, + "sample_count": 2024, + "synchronized_e2e_ms": { + "max": 40.66705, + "median": 0.191296, + "min": 0.168416, + "p90": 0.2273408, + "sample_count": 2024 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 2893, + "candidate_precomputed_gpu_span_ms": 0.030976, + "candidate_precomputed_host_enqueue_ms": 0.054464, + "candidate_precomputed_inter_kernel_gap_ms": 0.002016, + "candidate_precomputed_kernel_sum_ms": 0.02896, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.054464, + "candidate_precomputed_synchronized_e2e_ms": 0.06688, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.02896 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.036128, + "synchronized_e2e_ms": 0.053984 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.030976 + }, + "host_enqueue_ms": { + "median": 0.054464 + }, + "inter_kernel_gap_ms": { + "median": 0.002016 + }, + "kernel_sum_ms": { + "median": 0.02896 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2893, + "submission_ms": { + "median": 0.054464 + }, + "synchronized_e2e_ms": { + "median": 0.06688 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7f305ac0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7f306750" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.5971074380165289, + "submission": 0.8566392479435958, + "synchronized_e2e": 1.4521531100478469 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2024, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.049472, + "candidate_public_raw_host_enqueue_ms": 0.046656, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.049728, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.046656, + "candidate_public_raw_synchronized_e2e_ms": 0.09712, + "candidate_public_raw_tflops_from_gpu_span": 133.91217593790427, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.049472 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.048736, + "synchronized_e2e_ms": 0.086272 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.049472 + }, + "host_enqueue_ms": { + "median": 0.046656 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.049728 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2024, + "submission_ms": { + "median": 0.046656 + }, + "synchronized_e2e_ms": { + "median": 0.09712 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.072163, + "after_init_synchronized_e2e_ms_per_call": 2.101027, + "including_init_host_enqueue_ms_per_call": 36.946247, + "including_init_synchronized_e2e_ms_per_call": 37.044103, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.24920670000000006, + "after_init_synchronized_e2e_ms_per_call": 0.2975107, + "including_init_host_enqueue_ms_per_call": 3.7366151000000003, + "including_init_synchronized_e2e_ms_per_call": 3.7918183, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06691107, + "after_init_synchronized_e2e_ms_per_call": 0.11715906999999999, + "including_init_host_enqueue_ms_per_call": 0.41565191, + "including_init_synchronized_e2e_ms_per_call": 0.46658983, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.048681507000000006, + "after_init_synchronized_e2e_ms_per_call": 0.099123907, + "including_init_host_enqueue_ms_per_call": 0.083555591, + "including_init_synchronized_e2e_ms_per_call": 0.134066983, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.072163, + "synchronized_e2e_ms": 2.101027, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.050336, + "median": 0.049472, + "min": 0.049088, + "p90": 0.049856, + "sample_count": 2024 + }, + "host_enqueue_ms": { + "max": 0.353376, + "median": 0.046656, + "min": 0.03776, + "p90": 0.06591040000000001, + "sample_count": 2024 + }, + "sample_count": 2024, + "synchronized_e2e_ms": { + "max": 0.41568, + "median": 0.09712, + "min": 0.089696, + "p90": 0.1146176, + "sample_count": 2024 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.062496, + "submission_ms": 0.062496, + "synchronized_e2e_ms": 0.112384 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.046752, + "submission_ms": 0.046752, + "synchronized_e2e_ms": 0.094657 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.178272, + "submission_ms": 0.178272, + "synchronized_e2e_ms": 0.216384 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 6.782055, + "submission_ms": 6.782055, + "synchronized_e2e_ms": 6.808455 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.065024, + "submission_ms": 0.065024, + "synchronized_e2e_ms": 0.079872 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.036128, + "submission_ms": 0.036128, + "synchronized_e2e_ms": 0.053984 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.112545, + "submission_ms": 1.112545, + "synchronized_e2e_ms": 1.136705 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.006753, + "submission_ms": 1.006753, + "synchronized_e2e_ms": 1.029377 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.048736, + "submission_ms": 0.048736, + "synchronized_e2e_ms": 0.086272 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.072163, + "submission_ms": 2.072163, + "synchronized_e2e_ms": 2.101027 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.493791, + "evolution_kernel_ms": 0.189088, + "evolution_speedup": 2.6114, + "evolution_tflops": 35.0361, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_8f09_d416_midk_tail_b3_n3456_k768_d416", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 9834, + "measurement_schedule_sha256": "dd0c3e8cd7579ba0d1713b23b4e89e574e9ae2fd87414625d26e0d62879c55d7", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 2893, + "public_pair_count": 2024, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 2893, + "baseline_public_raw": 2024, + "candidate_precomputed": 2893, + "candidate_public_raw": 2024 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 1968 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:adjacent_8f09_d416_midk_tail_b3_n3456_k768_d416", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.6714876033057853, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.969686985172982, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.240536651837411, + "including_init_synchronized_e2e_speedup": 13.250420694489485, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.8671637692358627, + "including_init_synchronized_e2e_speedup": 12.990380182510325, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.197589909172205, + "including_init_synchronized_e2e_speedup": 10.925832931249273, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.9966238719787348, + "including_init_synchronized_e2e_speedup": 5.086671138113103, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.8211614650711512, + "hot_synchronized_e2e_speedup": 1.969686985172982, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 8094161, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_8f09_d416_midk_tail_b3_n3456_k768_d416", + "source": "tail_divisibility", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 480, + "K": 4096, + "N": 640, + "baseline_07cf_adapter_bench_iters": 1659, + "baseline_07cf_adapter_gpu_span_ms": 0.156096, + "baseline_07cf_adapter_host_enqueue_ms": 0.14864, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.027328, + "baseline_07cf_adapter_kernel_sum_ms": 0.128801, + "baseline_07cf_adapter_submission_ms": 0.14864, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.255968, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.128801 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.169024, + "synchronized_e2e_ms": 0.268416 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.156096 + }, + "host_enqueue_ms": { + "median": 0.14864 + }, + "inter_kernel_gap_ms": { + "median": 0.027328 + }, + "kernel_sum_ms": { + "median": 0.128801 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1659, + "submission_ms": { + "median": 0.14864 + }, + "synchronized_e2e_ms": { + "median": 0.255968 + } + }, + "baseline_07cf_precomputed_bench_iters": 2031, + "baseline_07cf_precomputed_gpu_span_ms": 0.128032, + "baseline_07cf_precomputed_host_enqueue_ms": 0.044096, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.128032, + "baseline_07cf_precomputed_submission_ms": 0.044096, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.177408, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.128032 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.048448, + "synchronized_e2e_ms": 0.15328 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.128032 + }, + "host_enqueue_ms": { + "median": 0.044096 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.128032 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2031, + "submission_ms": { + "median": 0.044096 + }, + "synchronized_e2e_ms": { + "median": 0.177408 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.2191952011997, + "submission": 3.3708272859216253, + "synchronized_e2e": 1.4428210678210676 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 6.9882, + "after_init_synchronized_e2e_ms_per_call": 7.075496, + "including_init_host_enqueue_ms_per_call": 42.539468, + "including_init_synchronized_e2e_ms_per_call": 42.706733, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.832596, + "after_init_synchronized_e2e_ms_per_call": 0.9379208, + "including_init_host_enqueue_ms_per_call": 4.387722800000001, + "including_init_synchronized_e2e_ms_per_call": 4.5010445, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2170356, + "after_init_synchronized_e2e_ms_per_call": 0.32416328, + "including_init_host_enqueue_ms_per_call": 0.5725482799999999, + "including_init_synchronized_e2e_ms_per_call": 0.6804756499999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15547956, + "after_init_synchronized_e2e_ms_per_call": 0.26278752799999994, + "including_init_host_enqueue_ms_per_call": 0.19103082799999999, + "including_init_synchronized_e2e_ms_per_call": 0.298418765, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 6.9882, + "synchronized_e2e_ms": 7.075496, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.299461, + "median": 0.156096, + "min": 0.146336, + "p90": 0.17509760000000002, + "sample_count": 1659 + }, + "host_enqueue_ms": { + "max": 26.498972, + "median": 0.14864, + "min": 0.124608, + "p90": 0.23008019999999998, + "sample_count": 1659 + }, + "sample_count": 1659, + "synchronized_e2e_ms": { + "max": 36.475494, + "median": 0.255968, + "min": 0.235136, + "p90": 0.32873600000000003, + "sample_count": 1659 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 2031, + "candidate_precomputed_gpu_span_ms": 0.049088, + "candidate_precomputed_host_enqueue_ms": 0.049984, + "candidate_precomputed_inter_kernel_gap_ms": 0.002112, + "candidate_precomputed_kernel_sum_ms": 0.04688, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.049984, + "candidate_precomputed_synchronized_e2e_ms": 0.083072, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.04688 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.036256, + "synchronized_e2e_ms": 0.063936 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.049088 + }, + "host_enqueue_ms": { + "median": 0.049984 + }, + "inter_kernel_gap_ms": { + "median": 0.002112 + }, + "kernel_sum_ms": { + "median": 0.04688 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2031, + "submission_ms": { + "median": 0.049984 + }, + "synchronized_e2e_ms": { + "median": 0.083072 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc027162a0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc02716240" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.2307692307692306, + "submission": 0.9507042253521126, + "synchronized_e2e": 1.307395993836672 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 1659, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.060416, + "candidate_public_raw_host_enqueue_ms": 0.04752, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.060352, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.04752, + "candidate_public_raw_synchronized_e2e_ms": 0.108608, + "candidate_public_raw_tflops_from_gpu_span": 83.3084745762712, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.060352 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.223968, + "synchronized_e2e_ms": 0.266496 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.060416 + }, + "host_enqueue_ms": { + "median": 0.04752 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.060352 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1659, + "submission_ms": { + "median": 0.04752 + }, + "synchronized_e2e_ms": { + "median": 0.108608 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.310306, + "after_init_synchronized_e2e_ms_per_call": 2.34445, + "including_init_host_enqueue_ms_per_call": 38.148326, + "including_init_synchronized_e2e_ms_per_call": 486.444663, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2737986, + "after_init_synchronized_e2e_ms_per_call": 0.33219220000000005, + "including_init_host_enqueue_ms_per_call": 3.8576006, + "including_init_synchronized_e2e_ms_per_call": 48.7422135, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07014786, + "after_init_synchronized_e2e_ms_per_call": 0.13096642, + "including_init_host_enqueue_ms_per_call": 0.42852806, + "including_init_synchronized_e2e_ms_per_call": 4.97196855, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.049782785999999996, + "after_init_synchronized_e2e_ms_per_call": 0.110843842, + "including_init_host_enqueue_ms_per_call": 0.085620806, + "including_init_synchronized_e2e_ms_per_call": 0.5949440549999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.310306, + "synchronized_e2e_ms": 2.34445, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.061504, + "median": 0.060416, + "min": 0.05888, + "p90": 0.060896, + "sample_count": 1659 + }, + "host_enqueue_ms": { + "max": 0.449952, + "median": 0.04752, + "min": 0.038304, + "p90": 0.072928, + "sample_count": 1659 + }, + "sample_count": 1659, + "synchronized_e2e_ms": { + "max": 0.500096, + "median": 0.108608, + "min": 0.100225, + "p90": 0.1286272, + "sample_count": 1659 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.143904, + "submission_ms": 0.143904, + "synchronized_e2e_ms": 0.24688 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.048448, + "submission_ms": 0.048448, + "synchronized_e2e_ms": 0.15328 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.169024, + "submission_ms": 0.169024, + "synchronized_e2e_ms": 0.268416 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 6.9882, + "submission_ms": 6.9882, + "synchronized_e2e_ms": 7.075496 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.059456, + "submission_ms": 0.059456, + "synchronized_e2e_ms": 0.07616 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.036256, + "submission_ms": 0.036256, + "synchronized_e2e_ms": 0.063936 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.237794, + "submission_ms": 1.237794, + "synchronized_e2e_ms": 1.25917 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.087873, + "submission_ms": 1.087873, + "synchronized_e2e_ms": 1.108769 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.223968, + "submission_ms": 0.223968, + "synchronized_e2e_ms": 0.266496 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.310306, + "submission_ms": 2.310306, + "synchronized_e2e_ms": 2.34445 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 1.999325, + "evolution_kernel_ms": 0.284192, + "evolution_speedup": 7.0351, + "evolution_tflops": 17.7104, + "expected_route": "d480_splitk_k1024_eac2_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_8f09_d480_highk_low_n_b2_n640_k4096_d480", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 7380, + "measurement_schedule_sha256": "d64a341c7952c085b52e2deacda6484b39e5ace1c8ac5f797be7244f75afb3cb", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 2031, + "public_pair_count": 1659, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 2031, + "baseline_public_raw": 1659, + "candidate_precomputed": 2031, + "candidate_public_raw": 1659 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 1478 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:adjacent_8f09_d480_highk_low_n_b2_n640_k4096_d480", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 2.608213820078227, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.3568061284619914, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.017976924225298, + "including_init_synchronized_e2e_speedup": 0.08779360993832099, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.8234281238391503, + "including_init_synchronized_e2e_speedup": 0.09234386739535332, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.4751633281264005, + "including_init_synchronized_e2e_speedup": 0.13686242041897065, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.370790503634834, + "including_init_synchronized_e2e_speedup": 0.5015913050849798, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.5836864406779663, + "hot_synchronized_e2e_speedup": 2.3568061284619914, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 8094801, + "selected_route": "d480_splitk_k1024_eac2_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_8f09_d480_highk_low_n_b2_n640_k4096_d480", + "source": "request_specific", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 48, + "K": 512, + "N": 1792, + "baseline_07cf_adapter_bench_iters": 7794, + "baseline_07cf_adapter_gpu_span_ms": 0.061441, + "baseline_07cf_adapter_host_enqueue_ms": 0.1669285, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.047968, + "baseline_07cf_adapter_kernel_sum_ms": 0.013472, + "baseline_07cf_adapter_submission_ms": 0.1669285, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.188032, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.013472 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.22864, + "synchronized_e2e_ms": 0.25088 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.061441 + }, + "host_enqueue_ms": { + "median": 0.1669285 + }, + "inter_kernel_gap_ms": { + "median": 0.047968 + }, + "kernel_sum_ms": { + "median": 0.013472 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 7794, + "submission_ms": { + "median": 0.1669285 + }, + "synchronized_e2e_ms": { + "median": 0.188032 + } + }, + "baseline_07cf_precomputed_bench_iters": 8955, + "baseline_07cf_precomputed_gpu_span_ms": 0.011168, + "baseline_07cf_precomputed_host_enqueue_ms": 0.052224, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.011168, + "baseline_07cf_precomputed_submission_ms": 0.052224, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.070016, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.011168 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.054528, + "synchronized_e2e_ms": 0.071104 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.011168 + }, + "host_enqueue_ms": { + "median": 0.052224 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.011168 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 8955, + "submission_ms": { + "median": 0.052224 + }, + "synchronized_e2e_ms": { + "median": 0.070016 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 5.501522206303725, + "submission": 3.196394378063726, + "synchronized_e2e": 2.6855575868372945 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.640328, + "after_init_synchronized_e2e_ms_per_call": 7.6682, + "including_init_host_enqueue_ms_per_call": 44.574766, + "including_init_synchronized_e2e_ms_per_call": 462.77951900000005, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.91426845, + "after_init_synchronized_e2e_ms_per_call": 0.9360488, + "including_init_host_enqueue_ms_per_call": 4.60771225, + "including_init_synchronized_e2e_ms_per_call": 46.447180700000004, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.241662495, + "after_init_synchronized_e2e_ms_per_call": 0.26283368, + "including_init_host_enqueue_ms_per_call": 0.611006875, + "including_init_synchronized_e2e_ms_per_call": 4.8139468700000005, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.1744018995, + "after_init_synchronized_e2e_ms_per_call": 0.19551216800000004, + "including_init_host_enqueue_ms_per_call": 0.21133633750000003, + "including_init_synchronized_e2e_ms_per_call": 0.6506234870000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.640328, + "synchronized_e2e_ms": 7.6682, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.990113, + "median": 0.061441, + "min": 0.049568, + "p90": 0.08498239999999999, + "sample_count": 7794 + }, + "host_enqueue_ms": { + "max": 88.365979, + "median": 0.1669285, + "min": 0.137728, + "p90": 0.2153738, + "sample_count": 7794 + }, + "sample_count": 7794, + "synchronized_e2e_ms": { + "max": 93.559041, + "median": 0.188032, + "min": 0.156, + "p90": 0.24152669999999998, + "sample_count": 7794 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 8955, + "candidate_precomputed_gpu_span_ms": 0.016576, + "candidate_precomputed_host_enqueue_ms": 0.063136, + "candidate_precomputed_inter_kernel_gap_ms": 0.007648, + "candidate_precomputed_kernel_sum_ms": 0.008928, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.063136, + "candidate_precomputed_synchronized_e2e_ms": 0.076064, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.008928 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.058848, + "synchronized_e2e_ms": 0.074592 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.016576 + }, + "host_enqueue_ms": { + "median": 0.063136 + }, + "inter_kernel_gap_ms": { + "median": 0.007648 + }, + "kernel_sum_ms": { + "median": 0.008928 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 8955, + "submission_ms": { + "median": 0.063136 + }, + "synchronized_e2e_ms": { + "median": 0.076064 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295c30e0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01aeda00" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 0.7953667953667953, + "submission": 0.9356310187531678, + "synchronized_e2e": 0.9854924800168279 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 7794, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.013184, + "candidate_public_raw_host_enqueue_ms": 0.059072, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.013376, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.059072, + "candidate_public_raw_synchronized_e2e_ms": 0.0749605, + "candidate_public_raw_tflops_from_gpu_span": 13.36170873786408, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.013184 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.069792, + "synchronized_e2e_ms": 0.089568 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.013184 + }, + "host_enqueue_ms": { + "median": 0.059072 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.013376 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 7794, + "submission_ms": { + "median": 0.059072 + }, + "synchronized_e2e_ms": { + "median": 0.0749605 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.989956, + "after_init_synchronized_e2e_ms_per_call": 3.018308, + "including_init_host_enqueue_ms_per_call": 40.347978999999995, + "including_init_synchronized_e2e_ms_per_call": 40.453931, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.3521604, + "after_init_synchronized_e2e_ms_per_call": 0.36929525, + "including_init_host_enqueue_ms_per_call": 4.087962699999999, + "including_init_synchronized_e2e_ms_per_call": 4.112857549999999, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.08838084, + "after_init_synchronized_e2e_ms_per_call": 0.104393975, + "including_init_host_enqueue_ms_per_call": 0.46196107, + "including_init_synchronized_e2e_ms_per_call": 0.478750205, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.062002884, + "after_init_synchronized_e2e_ms_per_call": 0.0779038475, + "including_init_host_enqueue_ms_per_call": 0.099360907, + "including_init_synchronized_e2e_ms_per_call": 0.11533947049999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.989956, + "synchronized_e2e_ms": 3.018308, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.013696, + "median": 0.013184, + "min": 0.012768, + "p90": 0.013408, + "sample_count": 7794 + }, + "host_enqueue_ms": { + "max": 92.028319, + "median": 0.059072, + "min": 0.043648, + "p90": 0.0792224, + "sample_count": 7794 + }, + "sample_count": 7794, + "synchronized_e2e_ms": { + "max": 92.316224, + "median": 0.0749605, + "min": 0.059872, + "p90": 0.100224, + "sample_count": 7794 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.076672, + "submission_ms": 0.076672, + "synchronized_e2e_ms": 0.095712 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.054528, + "submission_ms": 0.054528, + "synchronized_e2e_ms": 0.071104 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.22864, + "submission_ms": 0.22864, + "synchronized_e2e_ms": 0.25088 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.640328, + "submission_ms": 7.640328, + "synchronized_e2e_ms": 7.6682 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.099937, + "submission_ms": 0.099937, + "synchronized_e2e_ms": 0.118721 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.058848, + "submission_ms": 0.058848, + "synchronized_e2e_ms": 0.074592 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.309057, + "submission_ms": 1.309057, + "synchronized_e2e_ms": 1.336353 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.123265, + "submission_ms": 1.123265, + "synchronized_e2e_ms": 1.148737 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.069792, + "submission_ms": 0.069792, + "synchronized_e2e_ms": 0.089568 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.989956, + "submission_ms": 2.989956, + "synchronized_e2e_ms": 3.018308 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.316848, + "evolution_kernel_ms": 0.169375, + "evolution_speedup": 1.8707, + "evolution_tflops": 1.0401, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_8f09_d48_medium_tail_b2_n1792_k512_d48", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 33498, + "measurement_schedule_sha256": "f6fdd25b93c0f3f6aeee0692ab72ef29845fe0dc3fab7d798dbf8c5dee038a78", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 8955, + "public_pair_count": 7794, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 8955, + "baseline_public_raw": 7794, + "candidate_precomputed": 8955, + "candidate_public_raw": 7794 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 6700 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:adjacent_8f09_d48_medium_tail_b2_n1792_k512_d48", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.6737451737451737, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.508414431600643, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.5405624608224207, + "including_init_synchronized_e2e_speedup": 11.439667482500035, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.5346895201062023, + "including_init_synchronized_e2e_speedup": 11.293165429471296, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.517709283509896, + "including_init_synchronized_e2e_speedup": 10.055237198279634, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.509659975394669, + "including_init_synchronized_e2e_speedup": 5.640943938614667, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 4.660270024271845, + "hot_synchronized_e2e_speedup": 2.508414431600643, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 8090481, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_8f09_d48_medium_tail_b2_n1792_k512_d48", + "source": "tail_divisibility", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 64, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 3, + "D": 112, + "K": 768, + "N": 3840, + "baseline_07cf_adapter_bench_iters": 3692, + "baseline_07cf_adapter_gpu_span_ms": 0.060576, + "baseline_07cf_adapter_host_enqueue_ms": 0.144288, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.03344, + "baseline_07cf_adapter_kernel_sum_ms": 0.027168, + "baseline_07cf_adapter_submission_ms": 0.144288, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.163552, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.027168 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.359201, + "synchronized_e2e_ms": 0.389313 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.060576 + }, + "host_enqueue_ms": { + "median": 0.144288 + }, + "inter_kernel_gap_ms": { + "median": 0.03344 + }, + "kernel_sum_ms": { + "median": 0.027168 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3692, + "submission_ms": { + "median": 0.144288 + }, + "synchronized_e2e_ms": { + "median": 0.163552 + } + }, + "baseline_07cf_precomputed_bench_iters": 5140, + "baseline_07cf_precomputed_gpu_span_ms": 0.01952, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042177, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.01952, + "baseline_07cf_precomputed_submission_ms": 0.042177, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.068448, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.01952 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.064032, + "synchronized_e2e_ms": 0.087776 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.01952 + }, + "host_enqueue_ms": { + "median": 0.042177 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.01952 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5140, + "submission_ms": { + "median": 0.042177 + }, + "synchronized_e2e_ms": { + "median": 0.068448 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 3.10327868852459, + "submission": 3.4210114517390995, + "synchronized_e2e": 2.3894343151005146 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.876232, + "after_init_synchronized_e2e_ms_per_call": 7.913192, + "including_init_host_enqueue_ms_per_call": 42.104139, + "including_init_synchronized_e2e_ms_per_call": 42.258635000000005, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9174823999999999, + "after_init_synchronized_e2e_ms_per_call": 0.9385159999999999, + "including_init_host_enqueue_ms_per_call": 4.3402731, + "including_init_synchronized_e2e_ms_per_call": 4.373060300000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22160744000000002, + "after_init_synchronized_e2e_ms_per_call": 0.2410484, + "including_init_host_enqueue_ms_per_call": 0.56388651, + "including_init_synchronized_e2e_ms_per_call": 0.58450283, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.152019944, + "after_init_synchronized_e2e_ms_per_call": 0.17130164000000003, + "including_init_host_enqueue_ms_per_call": 0.186247851, + "including_init_synchronized_e2e_ms_per_call": 0.205647083, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.876232, + "synchronized_e2e_ms": 7.913192, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.321248, + "median": 0.060576, + "min": 0.051456, + "p90": 0.08409690000000002, + "sample_count": 3692 + }, + "host_enqueue_ms": { + "max": 54.271032, + "median": 0.144288, + "min": 0.122016, + "p90": 0.22434980000000002, + "sample_count": 3692 + }, + "sample_count": 3692, + "synchronized_e2e_ms": { + "max": 54.373304, + "median": 0.163552, + "min": 0.139776, + "p90": 0.2495168, + "sample_count": 3692 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 5140, + "candidate_precomputed_gpu_span_ms": 0.020832, + "candidate_precomputed_host_enqueue_ms": 0.053344, + "candidate_precomputed_inter_kernel_gap_ms": 0.004128, + "candidate_precomputed_kernel_sum_ms": 0.016704, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.053344, + "candidate_precomputed_synchronized_e2e_ms": 0.065344, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.016704 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.06272, + "synchronized_e2e_ms": 0.085536 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.020832 + }, + "host_enqueue_ms": { + "median": 0.053344 + }, + "inter_kernel_gap_ms": { + "median": 0.004128 + }, + "kernel_sum_ms": { + "median": 0.016704 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5140, + "submission_ms": { + "median": 0.053344 + }, + "synchronized_e2e_ms": { + "median": 0.065344 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282c6600", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282c4e30" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.3010752688172043, + "submission": 0.8233353329334133, + "synchronized_e2e": 1.1180215475024484 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 3692, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.027104, + "candidate_public_raw_host_enqueue_ms": 0.04392, + "candidate_public_raw_inter_kernel_gap_ms": 0.000128, + "candidate_public_raw_kernel_sum_ms": 0.026976, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.04392, + "candidate_public_raw_synchronized_e2e_ms": 0.073056, + "candidate_public_raw_tflops_from_gpu_span": 73.1186776859504, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.026976 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.114848, + "synchronized_e2e_ms": 0.145216 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.027104 + }, + "host_enqueue_ms": { + "median": 0.04392 + }, + "inter_kernel_gap_ms": { + "median": 0.000128 + }, + "kernel_sum_ms": { + "median": 0.026976 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3692, + "submission_ms": { + "median": 0.04392 + }, + "synchronized_e2e_ms": { + "median": 0.073056 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 51.719158, + "after_init_synchronized_e2e_ms_per_call": 51.771094, + "including_init_host_enqueue_ms_per_call": 86.29669799999999, + "including_init_synchronized_e2e_ms_per_call": 499.858886, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 5.2114438, + "after_init_synchronized_e2e_ms_per_call": 5.2428598, + "including_init_host_enqueue_ms_per_call": 8.6691978, + "including_init_synchronized_e2e_ms_per_call": 50.051639, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.5606723800000001, + "after_init_synchronized_e2e_ms_per_call": 0.5900363799999999, + "including_init_host_enqueue_ms_per_call": 0.9064477799999999, + "including_init_synchronized_e2e_ms_per_call": 5.0709143, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.095595238, + "after_init_synchronized_e2e_ms_per_call": 0.12475403799999998, + "including_init_host_enqueue_ms_per_call": 0.130172778, + "including_init_synchronized_e2e_ms_per_call": 0.57284183, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 51.719158, + "synchronized_e2e_ms": 51.771094, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.027872, + "median": 0.027104, + "min": 0.026752, + "p90": 0.02736, + "sample_count": 3692 + }, + "host_enqueue_ms": { + "max": 25.599994, + "median": 0.04392, + "min": 0.03472, + "p90": 0.06939200000000001, + "sample_count": 3692 + }, + "sample_count": 3692, + "synchronized_e2e_ms": { + "max": 50.122004, + "median": 0.073056, + "min": 0.065504, + "p90": 0.09369050000000004, + "sample_count": 3692 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.097248, + "submission_ms": 0.097248, + "synchronized_e2e_ms": 0.124992 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.064032, + "submission_ms": 0.064032, + "synchronized_e2e_ms": 0.087776 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.359201, + "submission_ms": 0.359201, + "synchronized_e2e_ms": 0.389313 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.876232, + "submission_ms": 7.876232, + "synchronized_e2e_ms": 7.913192 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.10288, + "submission_ms": 0.10288, + "synchronized_e2e_ms": 0.12896 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.06272, + "submission_ms": 0.06272, + "synchronized_e2e_ms": 0.085536 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.620801, + "submission_ms": 1.620801, + "synchronized_e2e_ms": 1.654946 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.438817, + "submission_ms": 1.438817, + "synchronized_e2e_ms": 1.472993 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.114848, + "submission_ms": 0.114848, + "synchronized_e2e_ms": 0.145216 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 51.719158, + "submission_ms": 51.719158, + "synchronized_e2e_ms": 51.771094 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.333232, + "evolution_kernel_ms": 0.194592, + "evolution_speedup": 1.7125, + "evolution_tflops": 10.1844, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_9ca0_d112_tail_div_b3_n3840_k768_d112", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 17664, + "measurement_schedule_sha256": "64aca44efa2373476270c5bcc8ea16e703b1fa0c879ebfd706fca32abbf61bd4", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 5140, + "public_pair_count": 3692, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 5140, + "baseline_public_raw": 3692, + "candidate_precomputed": 5140, + "candidate_public_raw": 3692 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 3534 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:adjacent_9ca0_d112_tail_div_b3_n3840_k768_d112", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.9370199692780338, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.238720981165134, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.1528496191330243, + "including_init_synchronized_e2e_speedup": 0.0845411298740021, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.17900841063878914, + "including_init_synchronized_e2e_speedup": 0.08737097100856178, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.40853141970669676, + "including_init_synchronized_e2e_speedup": 0.11526576775316436, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.3731149928790285, + "including_init_synchronized_e2e_speedup": 0.3589945290831852, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.2349468713105076, + "hot_synchronized_e2e_speedup": 2.238720981165134, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 921125, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_9ca0_d112_tail_div_b3_n3840_k768_d112", + "source": "tail_divisibility", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 8, + "D": 128, + "K": 256, + "N": 8064, + "baseline_07cf_adapter_bench_iters": 1811, + "baseline_07cf_adapter_gpu_span_ms": 0.067297, + "baseline_07cf_adapter_host_enqueue_ms": 0.164512, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.011296, + "baseline_07cf_adapter_kernel_sum_ms": 0.056032, + "baseline_07cf_adapter_submission_ms": 0.164512, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.184928, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.056032 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.263457, + "synchronized_e2e_ms": 0.288993 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.067297 + }, + "host_enqueue_ms": { + "median": 0.164512 + }, + "inter_kernel_gap_ms": { + "median": 0.011296 + }, + "kernel_sum_ms": { + "median": 0.056032 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1811, + "submission_ms": { + "median": 0.164512 + }, + "synchronized_e2e_ms": { + "median": 0.184928 + } + }, + "baseline_07cf_precomputed_bench_iters": 5601, + "baseline_07cf_precomputed_gpu_span_ms": 0.018336, + "baseline_07cf_precomputed_host_enqueue_ms": 0.045056, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.018336, + "baseline_07cf_precomputed_submission_ms": 0.045056, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.068768, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.018336 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.065376, + "synchronized_e2e_ms": 0.086688 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.018336 + }, + "host_enqueue_ms": { + "median": 0.045056 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.018336 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5601, + "submission_ms": { + "median": 0.045056 + }, + "synchronized_e2e_ms": { + "median": 0.068768 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 3.6702116055846417, + "submission": 3.651278409090909, + "synchronized_e2e": 2.68915774778967 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.753704, + "after_init_synchronized_e2e_ms_per_call": 7.787176, + "including_init_host_enqueue_ms_per_call": 42.203148, + "including_init_synchronized_e2e_ms_per_call": 491.82867, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9234311999999999, + "after_init_synchronized_e2e_ms_per_call": 0.9451528, + "including_init_host_enqueue_ms_per_call": 4.368375599999999, + "including_init_synchronized_e2e_ms_per_call": 49.3493022, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.24040391999999997, + "after_init_synchronized_e2e_ms_per_call": 0.26095048, + "including_init_host_enqueue_ms_per_call": 0.58489836, + "including_init_synchronized_e2e_ms_per_call": 5.10136542, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.172101192, + "after_init_synchronized_e2e_ms_per_call": 0.192530248, + "including_init_host_enqueue_ms_per_call": 0.206550636, + "including_init_synchronized_e2e_ms_per_call": 0.6765717419999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.753704, + "synchronized_e2e_ms": 7.787176, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.318881, + "median": 0.067297, + "min": 0.056672, + "p90": 0.09168, + "sample_count": 1811 + }, + "host_enqueue_ms": { + "max": 27.106748, + "median": 0.164512, + "min": 0.13184, + "p90": 0.216288, + "sample_count": 1811 + }, + "sample_count": 1811, + "synchronized_e2e_ms": { + "max": 27.148124, + "median": 0.184928, + "min": 0.150784, + "p90": 0.240864, + "sample_count": 1811 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 5601, + "candidate_precomputed_gpu_span_ms": 0.017728, + "candidate_precomputed_host_enqueue_ms": 0.041888, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.017728, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.041888, + "candidate_precomputed_synchronized_e2e_ms": 0.053984, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.017728 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.05776, + "synchronized_e2e_ms": 0.07568 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.017728 + }, + "host_enqueue_ms": { + "median": 0.041888 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.017728 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5601, + "submission_ms": { + "median": 0.041888 + }, + "synchronized_e2e_ms": { + "median": 0.053984 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7edfdf10", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7edff8f0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 3.111913357400722, + "submission": 1.3063407181054238, + "synchronized_e2e": 2.036751630112626 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 1811, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.055168, + "candidate_public_raw_host_enqueue_ms": 0.05472, + "candidate_public_raw_inter_kernel_gap_ms": 0.00016, + "candidate_public_raw_kernel_sum_ms": 0.055008, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.05472, + "candidate_public_raw_synchronized_e2e_ms": 0.109952, + "candidate_public_raw_tflops_from_gpu_span": 76.63606496519722, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.055008 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.328352, + "synchronized_e2e_ms": 0.369248 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.055168 + }, + "host_enqueue_ms": { + "median": 0.05472 + }, + "inter_kernel_gap_ms": { + "median": 0.00016 + }, + "kernel_sum_ms": { + "median": 0.055008 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1811, + "submission_ms": { + "median": 0.05472 + }, + "synchronized_e2e_ms": { + "median": 0.109952 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.557411, + "after_init_synchronized_e2e_ms_per_call": 2.589123, + "including_init_host_enqueue_ms_per_call": 37.431495, + "including_init_synchronized_e2e_ms_per_call": 37.532199, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.3049891, + "after_init_synchronized_e2e_ms_per_call": 0.3578691, + "including_init_host_enqueue_ms_per_call": 3.7923975, + "including_init_synchronized_e2e_ms_per_call": 3.8521767, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07974691, + "after_init_synchronized_e2e_ms_per_call": 0.13474371, + "including_init_host_enqueue_ms_per_call": 0.42848774999999995, + "including_init_synchronized_e2e_ms_per_call": 0.48417446999999997, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.057222691, + "after_init_synchronized_e2e_ms_per_call": 0.112431171, + "including_init_host_enqueue_ms_per_call": 0.09209677499999999, + "including_init_synchronized_e2e_ms_per_call": 0.14737424699999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.557411, + "synchronized_e2e_ms": 2.589123, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.056064, + "median": 0.055168, + "min": 0.054432, + "p90": 0.05552, + "sample_count": 1811 + }, + "host_enqueue_ms": { + "max": 5.477766, + "median": 0.05472, + "min": 0.038784, + "p90": 0.07616, + "sample_count": 1811 + }, + "sample_count": 1811, + "synchronized_e2e_ms": { + "max": 15.767025, + "median": 0.109952, + "min": 0.095424, + "p90": 0.127712, + "sample_count": 1811 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.088128, + "submission_ms": 0.088128, + "synchronized_e2e_ms": 0.109856 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.065376, + "submission_ms": 0.065376, + "synchronized_e2e_ms": 0.086688 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.263457, + "submission_ms": 0.263457, + "synchronized_e2e_ms": 0.288993 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.753704, + "submission_ms": 7.753704, + "synchronized_e2e_ms": 7.787176 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.083169, + "submission_ms": 0.083169, + "synchronized_e2e_ms": 0.105601 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.05776, + "submission_ms": 0.05776, + "synchronized_e2e_ms": 0.07568 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.416385, + "submission_ms": 1.416385, + "synchronized_e2e_ms": 1.443937 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.265857, + "submission_ms": 1.265857, + "synchronized_e2e_ms": 1.291745 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.328352, + "submission_ms": 0.328352, + "synchronized_e2e_ms": 0.369248 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.557411, + "submission_ms": 2.557411, + "synchronized_e2e_ms": 2.589123 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.291264, + "evolution_kernel_ms": 0.155168, + "evolution_speedup": 1.8771, + "evolution_tflops": 27.247, + "expected_route": "aligned_weave_v10_fallback", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_9ca0_d128_exact_guard_neighbor_b8_n8064_k256_d128", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 14824, + "measurement_schedule_sha256": "05024f867dd45e2599655d22aa298623e85c577053441328d4164ee44eb5d7c0", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 5601, + "public_pair_count": 1811, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 5601, + "baseline_public_raw": 1811, + "candidate_precomputed": 5601, + "candidate_public_raw": 1811 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2968 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:adjacent_9ca0_d128_exact_guard_neighbor_b8_n8064_k256_d128", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.0342960288808665, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.681897555296857, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.0076500807416258, + "including_init_synchronized_e2e_speedup": 13.104179427376478, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.6410573027959106, + "including_init_synchronized_e2e_speedup": 12.81075766851505, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 1.9366431279055625, + "including_init_synchronized_e2e_speedup": 10.536213154733252, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.7124276683020583, + "including_init_synchronized_e2e_speedup": 4.590841044297244, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.2198557134570764, + "hot_synchronized_e2e_speedup": 1.681897555296857, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 912806, + "selected_route": "aligned_weave_v10_fallback", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_9ca0_d128_exact_guard_neighbor_b8_n8064_k256_d128", + "source": "guard_miss_fallback", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 224, + "K": 256, + "N": 4096, + "baseline_07cf_adapter_bench_iters": 4377, + "baseline_07cf_adapter_gpu_span_ms": 0.054208, + "baseline_07cf_adapter_host_enqueue_ms": 0.157345, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.033472, + "baseline_07cf_adapter_kernel_sum_ms": 0.020704, + "baseline_07cf_adapter_submission_ms": 0.157345, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.17648, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.020704 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.226144, + "synchronized_e2e_ms": 0.25168 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.054208 + }, + "host_enqueue_ms": { + "median": 0.157345 + }, + "inter_kernel_gap_ms": { + "median": 0.033472 + }, + "kernel_sum_ms": { + "median": 0.020704 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 4377, + "submission_ms": { + "median": 0.157345 + }, + "synchronized_e2e_ms": { + "median": 0.17648 + } + }, + "baseline_07cf_precomputed_bench_iters": 9890, + "baseline_07cf_precomputed_gpu_span_ms": 0.010048, + "baseline_07cf_precomputed_host_enqueue_ms": 0.043424, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.010048, + "baseline_07cf_precomputed_submission_ms": 0.043424, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.058848, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.010048 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.044065, + "synchronized_e2e_ms": 0.060449 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.010048 + }, + "host_enqueue_ms": { + "median": 0.043424 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.010048 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 9890, + "submission_ms": { + "median": 0.043424 + }, + "synchronized_e2e_ms": { + "median": 0.058848 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 5.3949044585987265, + "submission": 3.6234570744288876, + "synchronized_e2e": 2.9989124524197934 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.284551, + "after_init_synchronized_e2e_ms_per_call": 7.312775, + "including_init_host_enqueue_ms_per_call": 42.835819, + "including_init_synchronized_e2e_ms_per_call": 42.944012, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8700656, + "after_init_synchronized_e2e_ms_per_call": 0.8901095, + "including_init_host_enqueue_ms_per_call": 4.4251924, + "including_init_synchronized_e2e_ms_per_call": 4.453233200000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22861706, + "after_init_synchronized_e2e_ms_per_call": 0.24784295, + "including_init_host_enqueue_ms_per_call": 0.5841297400000001, + "including_init_synchronized_e2e_ms_per_call": 0.6041553199999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.164472206, + "after_init_synchronized_e2e_ms_per_call": 0.18361629499999998, + "including_init_host_enqueue_ms_per_call": 0.20002347400000003, + "including_init_synchronized_e2e_ms_per_call": 0.21924753199999997, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.284551, + "synchronized_e2e_ms": 7.312775, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.394081, + "median": 0.054208, + "min": 0.04256, + "p90": 0.0795968, + "sample_count": 4377 + }, + "host_enqueue_ms": { + "max": 0.977857, + "median": 0.157345, + "min": 0.125856, + "p90": 0.2574080000000001, + "sample_count": 4377 + }, + "sample_count": 4377, + "synchronized_e2e_ms": { + "max": 5.717254, + "median": 0.17648, + "min": 0.14112, + "p90": 0.2854342, + "sample_count": 4377 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 9890, + "candidate_precomputed_gpu_span_ms": 0.0144, + "candidate_precomputed_host_enqueue_ms": 0.053711999999999996, + "candidate_precomputed_inter_kernel_gap_ms": 0.003616, + "candidate_precomputed_kernel_sum_ms": 0.010816, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.053711999999999996, + "candidate_precomputed_synchronized_e2e_ms": 0.065696, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.010816 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.043616, + "synchronized_e2e_ms": 0.059104 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.0144 + }, + "host_enqueue_ms": { + "median": 0.053711999999999996 + }, + "inter_kernel_gap_ms": { + "median": 0.003616 + }, + "kernel_sum_ms": { + "median": 0.010816 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 9890, + "submission_ms": { + "median": 0.053711999999999996 + }, + "synchronized_e2e_ms": { + "median": 0.065696 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc0121fb90", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc0121e990" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.6022222222222222, + "submission": 0.8632707774798928, + "synchronized_e2e": 1.0662445202143205 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 4377, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.023072, + "candidate_public_raw_host_enqueue_ms": 0.046368, + "candidate_public_raw_inter_kernel_gap_ms": 3.2e-05, + "candidate_public_raw_kernel_sum_ms": 0.023168, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.046368, + "candidate_public_raw_synchronized_e2e_ms": 0.070048, + "candidate_public_raw_tflops_from_gpu_span": 40.721398058252426, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.02304 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.070976, + "synchronized_e2e_ms": 0.090881 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.023072 + }, + "host_enqueue_ms": { + "median": 0.046368 + }, + "inter_kernel_gap_ms": { + "median": 3.2e-05 + }, + "kernel_sum_ms": { + "median": 0.023168 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 4377, + "submission_ms": { + "median": 0.046368 + }, + "synchronized_e2e_ms": { + "median": 0.070048 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.206179, + "after_init_synchronized_e2e_ms_per_call": 2.232675, + "including_init_host_enqueue_ms_per_call": 38.044199, + "including_init_synchronized_e2e_ms_per_call": 486.33288799999997, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2623491, + "after_init_synchronized_e2e_ms_per_call": 0.2863107, + "including_init_host_enqueue_ms_per_call": 3.8461511, + "including_init_synchronized_e2e_ms_per_call": 48.696332, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06796611000000001, + "after_init_synchronized_e2e_ms_per_call": 0.09167427, + "including_init_host_enqueue_ms_per_call": 0.42634631, + "including_init_synchronized_e2e_ms_per_call": 4.9326764, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.048527811, + "after_init_synchronized_e2e_ms_per_call": 0.072210627, + "including_init_host_enqueue_ms_per_call": 0.084365831, + "including_init_synchronized_e2e_ms_per_call": 0.55631084, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.206179, + "synchronized_e2e_ms": 2.232675, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.02384, + "median": 0.023072, + "min": 0.022528, + "p90": 0.02336, + "sample_count": 4377 + }, + "host_enqueue_ms": { + "max": 31.162688, + "median": 0.046368, + "min": 0.037024, + "p90": 0.080192, + "sample_count": 4377 + }, + "sample_count": 4377, + "synchronized_e2e_ms": { + "max": 55.000729, + "median": 0.070048, + "min": 0.061984, + "p90": 0.10354560000000003, + "sample_count": 4377 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.063296, + "submission_ms": 0.063296, + "synchronized_e2e_ms": 0.082048 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.044065, + "submission_ms": 0.044065, + "synchronized_e2e_ms": 0.060449 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.226144, + "submission_ms": 0.226144, + "synchronized_e2e_ms": 0.25168 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.284551, + "submission_ms": 7.284551, + "synchronized_e2e_ms": 7.312775 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.068928, + "submission_ms": 0.068928, + "synchronized_e2e_ms": 0.08656 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.043616, + "submission_ms": 0.043616, + "synchronized_e2e_ms": 0.059104 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.219137, + "submission_ms": 1.219137, + "synchronized_e2e_ms": 1.243329 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.089313, + "submission_ms": 1.089313, + "synchronized_e2e_ms": 1.114145 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.070976, + "submission_ms": 0.070976, + "synchronized_e2e_ms": 0.090881 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.206179, + "submission_ms": 2.206179, + "synchronized_e2e_ms": 2.232675 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.258143, + "evolution_kernel_ms": 0.1712, + "evolution_speedup": 1.5078, + "evolution_tflops": 5.4879, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_9ca0_d224_guard_overlap_b2_n4096_k256_d224", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 28534, + "measurement_schedule_sha256": "cc30ec37651ba653c0421bada2a72f98e4cfbe9e99d92c84412dd3450669c7aa", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 9890, + "public_pair_count": 4377, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 9890, + "baseline_public_raw": 4377, + "candidate_precomputed": 9890, + "candidate_public_raw": 4377 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 5708 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:adjacent_9ca0_d224_guard_overlap_b2_n4096_k256_d224", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.6977777777777777, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.5194152581087255, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.275342358381762, + "including_init_synchronized_e2e_speedup": 0.08830168195411042, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 3.1088935900753976, + "including_init_synchronized_e2e_speedup": 0.09144904794882705, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.703517028278491, + "including_init_synchronized_e2e_speedup": 0.12248022594792554, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.542787711841915, + "including_init_synchronized_e2e_speedup": 0.39410976065107767, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.3495145631067964, + "hot_synchronized_e2e_speedup": 2.5194152581087255, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 922405, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_9ca0_d224_guard_overlap_b2_n4096_k256_d224", + "source": "guard_overlap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 288, + "K": 4096, + "N": 384, + "baseline_07cf_adapter_bench_iters": 3692, + "baseline_07cf_adapter_gpu_span_ms": 0.154576, + "baseline_07cf_adapter_host_enqueue_ms": 0.16835250000000002, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.044736, + "baseline_07cf_adapter_kernel_sum_ms": 0.109824, + "baseline_07cf_adapter_submission_ms": 0.16835250000000002, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.2677445, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.109824 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.265568, + "synchronized_e2e_ms": 0.360353 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.154576 + }, + "host_enqueue_ms": { + "median": 0.16835250000000002 + }, + "inter_kernel_gap_ms": { + "median": 0.044736 + }, + "kernel_sum_ms": { + "median": 0.109824 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3692, + "submission_ms": { + "median": 0.16835250000000002 + }, + "synchronized_e2e_ms": { + "median": 0.2677445 + } + }, + "baseline_07cf_precomputed_bench_iters": 3502, + "baseline_07cf_precomputed_gpu_span_ms": 0.121632, + "baseline_07cf_precomputed_host_enqueue_ms": 0.053056, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.121632, + "baseline_07cf_precomputed_submission_ms": 0.053056, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.1814405, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.121632 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.052608, + "synchronized_e2e_ms": 0.151488 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.121632 + }, + "host_enqueue_ms": { + "median": 0.053056 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.121632 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3502, + "submission_ms": { + "median": 0.053056 + }, + "synchronized_e2e_ms": { + "median": 0.1814405 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.270849776374638, + "submission": 3.173109544632087, + "synchronized_e2e": 1.4756600648697507 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.724904, + "after_init_synchronized_e2e_ms_per_call": 7.787688, + "including_init_host_enqueue_ms_per_call": 44.659342, + "including_init_synchronized_e2e_ms_per_call": 462.89900700000004, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9240076500000001, + "after_init_synchronized_e2e_ms_per_call": 1.01973885, + "including_init_host_enqueue_ms_per_call": 4.61745145, + "including_init_synchronized_e2e_ms_per_call": 46.530870750000005, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.243918015, + "after_init_synchronized_e2e_ms_per_call": 0.342943935, + "including_init_host_enqueue_ms_per_call": 0.613262395, + "including_init_synchronized_e2e_ms_per_call": 4.894057125000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.17590905150000002, + "after_init_synchronized_e2e_ms_per_call": 0.2752644435, + "including_init_host_enqueue_ms_per_call": 0.21284348950000004, + "including_init_synchronized_e2e_ms_per_call": 0.7303757625, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.724904, + "synchronized_e2e_ms": 7.787688, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.35639, + "median": 0.154576, + "min": 0.142625, + "p90": 0.177857, + "sample_count": 3692 + }, + "host_enqueue_ms": { + "max": 33.32509, + "median": 0.16835250000000002, + "min": 0.13712, + "p90": 0.2075904, + "sample_count": 3692 + }, + "sample_count": 3692, + "synchronized_e2e_ms": { + "max": 82.932598, + "median": 0.2677445, + "min": 0.238528, + "p90": 0.30301130000000004, + "sample_count": 3692 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3502, + "candidate_precomputed_gpu_span_ms": 0.020512, + "candidate_precomputed_host_enqueue_ms": 0.059264, + "candidate_precomputed_inter_kernel_gap_ms": 0.002208, + "candidate_precomputed_kernel_sum_ms": 0.018288, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.059264, + "candidate_precomputed_synchronized_e2e_ms": 0.07184, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.018288 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.047936, + "synchronized_e2e_ms": 0.062048 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.020512 + }, + "host_enqueue_ms": { + "median": 0.059264 + }, + "inter_kernel_gap_ms": { + "median": 0.002208 + }, + "kernel_sum_ms": { + "median": 0.018288 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3502, + "submission_ms": { + "median": 0.059264 + }, + "synchronized_e2e_ms": { + "median": 0.07184 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc281f9460", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc281faa80" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.3463338533541342, + "submission": 1.039416846652268, + "synchronized_e2e": 1.2685968819599107 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 3692, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.027616, + "candidate_public_raw_host_enqueue_ms": 0.0616, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.027648, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.0616, + "candidate_public_raw_synchronized_e2e_ms": 0.091136, + "candidate_public_raw_tflops_from_gpu_span": 32.80596987253766, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.027552 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.266624, + "synchronized_e2e_ms": 0.288512 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.027616 + }, + "host_enqueue_ms": { + "median": 0.0616 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.027648 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3692, + "submission_ms": { + "median": 0.0616 + }, + "synchronized_e2e_ms": { + "median": 0.091136 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 67.790118, + "after_init_synchronized_e2e_ms_per_call": 67.819078, + "including_init_host_enqueue_ms_per_call": 105.14814100000001, + "including_init_synchronized_e2e_ms_per_call": 105.25470100000001, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 6.834451800000001, + "after_init_synchronized_e2e_ms_per_call": 6.8639302, + "including_init_host_enqueue_ms_per_call": 10.570254100000001, + "including_init_synchronized_e2e_ms_per_call": 10.607492500000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.73888518, + "after_init_synchronized_e2e_ms_per_call": 0.7684154200000001, + "including_init_host_enqueue_ms_per_call": 1.11246541, + "including_init_synchronized_e2e_ms_per_call": 1.14277165, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.129328518, + "after_init_synchronized_e2e_ms_per_call": 0.158863942, + "including_init_host_enqueue_ms_per_call": 0.166686541, + "including_init_synchronized_e2e_ms_per_call": 0.196299565, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 67.790118, + "synchronized_e2e_ms": 67.819078, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.028512, + "median": 0.027616, + "min": 0.02656, + "p90": 0.028032, + "sample_count": 3692 + }, + "host_enqueue_ms": { + "max": 25.673531, + "median": 0.0616, + "min": 0.041408, + "p90": 0.0794848, + "sample_count": 3692 + }, + "sample_count": 3692, + "synchronized_e2e_ms": { + "max": 34.858052, + "median": 0.091136, + "min": 0.072768, + "p90": 0.1059785, + "sample_count": 3692 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.08096, + "submission_ms": 0.08096, + "synchronized_e2e_ms": 0.178496 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.052608, + "submission_ms": 0.052608, + "synchronized_e2e_ms": 0.151488 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.265568, + "submission_ms": 0.265568, + "synchronized_e2e_ms": 0.360353 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.724904, + "submission_ms": 7.724904, + "synchronized_e2e_ms": 7.787688 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.087904, + "submission_ms": 0.087904, + "synchronized_e2e_ms": 0.10464 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.047936, + "submission_ms": 0.047936, + "synchronized_e2e_ms": 0.062048 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.354657, + "submission_ms": 1.354657, + "synchronized_e2e_ms": 1.379361 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.164001, + "submission_ms": 1.164001, + "synchronized_e2e_ms": 1.18653 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.266624, + "submission_ms": 0.266624, + "synchronized_e2e_ms": 0.288512 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 67.790118, + "submission_ms": 67.790118, + "synchronized_e2e_ms": 67.819078 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 56.807529, + "evolution_kernel_ms": 0.243872, + "evolution_speedup": 232.9399, + "evolution_tflops": 3.7149, + "expected_route": "d288_parent_splitk_hybrid_20260629_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_9ca0_d288_highk_low_n_b1_n384_k4096_d288", + "measurement_order": [ + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 14388, + "measurement_schedule_sha256": "53fa22f5f33946f66c72dfad7b2a39b51ba34a00622c4fde2b122033e2e8a01a", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3502, + "public_pair_count": 3692, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3502, + "baseline_public_raw": 3692, + "candidate_precomputed": 3502, + "candidate_public_raw": 3692 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2880 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:adjacent_9ca0_d288_highk_low_n_b1_n384_k4096_d288", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 5.929797191887676, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.9378566099016856, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.11483034316685932, + "including_init_synchronized_e2e_speedup": 4.397893895494511, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.1485648630284731, + "including_init_synchronized_e2e_speedup": 4.386604162105229, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.44630017315373494, + "including_init_synchronized_e2e_speedup": 4.282620351143644, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.7327056098104376, + "including_init_synchronized_e2e_speedup": 3.72072022930871, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 5.597334878331401, + "hot_synchronized_e2e_speedup": 2.9378566099016856, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 928805, + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_9ca0_d288_highk_low_n_b1_n384_k4096_d288", + "source": "heldout_neighborhood", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 352, + "K": 8192, + "N": 512, + "baseline_07cf_adapter_bench_iters": 2327, + "baseline_07cf_adapter_gpu_span_ms": 0.250336, + "baseline_07cf_adapter_host_enqueue_ms": 0.151904, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.029952, + "baseline_07cf_adapter_kernel_sum_ms": 0.220384, + "baseline_07cf_adapter_submission_ms": 0.151904, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.355136, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.220384 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.195488, + "synchronized_e2e_ms": 0.383745 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.250336 + }, + "host_enqueue_ms": { + "median": 0.151904 + }, + "inter_kernel_gap_ms": { + "median": 0.029952 + }, + "kernel_sum_ms": { + "median": 0.220384 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2327, + "submission_ms": { + "median": 0.151904 + }, + "synchronized_e2e_ms": { + "median": 0.355136 + } + }, + "baseline_07cf_precomputed_bench_iters": 3242, + "baseline_07cf_precomputed_gpu_span_ms": 0.246496, + "baseline_07cf_precomputed_host_enqueue_ms": 0.044064, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.246496, + "baseline_07cf_precomputed_submission_ms": 0.044064, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.2976, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.246496 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.044, + "synchronized_e2e_ms": 0.232864 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.246496 + }, + "host_enqueue_ms": { + "median": 0.044064 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.246496 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3242, + "submission_ms": { + "median": 0.044064 + }, + "synchronized_e2e_ms": { + "median": 0.2976 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.0155783460989225, + "submission": 3.4473493100944084, + "synchronized_e2e": 1.1933333333333334 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.313128, + "after_init_synchronized_e2e_ms_per_call": 7.487496, + "including_init_host_enqueue_ms_per_call": 41.541035, + "including_init_synchronized_e2e_ms_per_call": 41.832939, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8680264, + "after_init_synchronized_e2e_ms_per_call": 1.068372, + "including_init_host_enqueue_ms_per_call": 4.2908171, + "including_init_synchronized_e2e_ms_per_call": 4.502916300000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22351624, + "after_init_synchronized_e2e_ms_per_call": 0.42645960000000005, + "including_init_host_enqueue_ms_per_call": 0.56579531, + "including_init_synchronized_e2e_ms_per_call": 0.76991403, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15906522400000003, + "after_init_synchronized_e2e_ms_per_call": 0.36226836, + "including_init_host_enqueue_ms_per_call": 0.193293131, + "including_init_synchronized_e2e_ms_per_call": 0.39661380300000004, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.313128, + "synchronized_e2e_ms": 7.487496, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.525185, + "median": 0.250336, + "min": 0.241856, + "p90": 0.270208, + "sample_count": 2327 + }, + "host_enqueue_ms": { + "max": 46.232015, + "median": 0.151904, + "min": 0.125792, + "p90": 0.23825320000000005, + "sample_count": 2327 + }, + "sample_count": 2327, + "synchronized_e2e_ms": { + "max": 46.378127, + "median": 0.355136, + "min": 0.330848, + "p90": 0.4301706, + "sample_count": 2327 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3242, + "candidate_precomputed_gpu_span_ms": 0.030656, + "candidate_precomputed_host_enqueue_ms": 0.051808, + "candidate_precomputed_inter_kernel_gap_ms": 0.002048, + "candidate_precomputed_kernel_sum_ms": 0.028576, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.051808, + "candidate_precomputed_synchronized_e2e_ms": 0.067136, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.028576 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.042561, + "synchronized_e2e_ms": 0.057025 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.030656 + }, + "host_enqueue_ms": { + "median": 0.051808 + }, + "inter_kernel_gap_ms": { + "median": 0.002048 + }, + "kernel_sum_ms": { + "median": 0.028576 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3242, + "submission_ms": { + "median": 0.051808 + }, + "synchronized_e2e_ms": { + "median": 0.067136 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb9e5ae6f0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb9e5afc50" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.407098121085595, + "submission": 0.9505867819641755, + "synchronized_e2e": 1.4013346043851287 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2327, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.043136, + "candidate_public_raw_host_enqueue_ms": 0.049248, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.043296, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.049248, + "candidate_public_raw_synchronized_e2e_ms": 0.09408, + "candidate_public_raw_tflops_from_gpu_span": 68.45303264094956, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.043136 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.293984, + "synchronized_e2e_ms": 0.320224 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.043136 + }, + "host_enqueue_ms": { + "median": 0.049248 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.043296 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2327, + "submission_ms": { + "median": 0.049248 + }, + "synchronized_e2e_ms": { + "median": 0.09408 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.674211, + "after_init_synchronized_e2e_ms_per_call": 2.698531, + "including_init_host_enqueue_ms_per_call": 37.251751, + "including_init_synchronized_e2e_ms_per_call": 450.786323, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.31174430000000003, + "after_init_synchronized_e2e_ms_per_call": 0.3545251, + "including_init_host_enqueue_ms_per_call": 3.7694983, + "including_init_synchronized_e2e_ms_per_call": 45.1633043, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07549763000000001, + "after_init_synchronized_e2e_ms_per_call": 0.12012450999999999, + "including_init_host_enqueue_ms_per_call": 0.42127303, + "including_init_synchronized_e2e_ms_per_call": 4.601002429999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.051872963, + "after_init_synchronized_e2e_ms_per_call": 0.09668445099999999, + "including_init_host_enqueue_ms_per_call": 0.086450503, + "including_init_synchronized_e2e_ms_per_call": 0.544772243, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.674211, + "synchronized_e2e_ms": 2.698531, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.04464, + "median": 0.043136, + "min": 0.041728, + "p90": 0.043904, + "sample_count": 2327 + }, + "host_enqueue_ms": { + "max": 0.597888, + "median": 0.049248, + "min": 0.038272, + "p90": 0.07800960000000004, + "sample_count": 2327 + }, + "sample_count": 2327, + "synchronized_e2e_ms": { + "max": 0.74576, + "median": 0.09408, + "min": 0.084288, + "p90": 0.11664740000000001, + "sample_count": 2327 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.067328, + "submission_ms": 0.067328, + "synchronized_e2e_ms": 0.254944 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.044, + "submission_ms": 0.044, + "synchronized_e2e_ms": 0.232864 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.195488, + "submission_ms": 0.195488, + "synchronized_e2e_ms": 0.383745 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.313128, + "submission_ms": 7.313128, + "synchronized_e2e_ms": 7.487496 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.067392, + "submission_ms": 0.067392, + "synchronized_e2e_ms": 0.082432 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.042561, + "submission_ms": 0.042561, + "synchronized_e2e_ms": 0.057025 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.296161, + "submission_ms": 1.296161, + "synchronized_e2e_ms": 1.316385 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.240673, + "submission_ms": 1.240673, + "synchronized_e2e_ms": 1.262209 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.293984, + "submission_ms": 0.293984, + "synchronized_e2e_ms": 0.320224 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.674211, + "submission_ms": 2.674211, + "synchronized_e2e_ms": 2.698531 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 49.145822, + "evolution_kernel_ms": 0.275519, + "evolution_speedup": 178.3754, + "evolution_tflops": 10.7172, + "expected_route": "d352_exactd_splitk_c95c_v2", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_9ca0_d352_splitk_boundary_b1_n512_k8192_d352", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 11138, + "measurement_schedule_sha256": "cbc7534d7a3ffbaaddeb176e6f1fa598077343dfa8367e455e4dbba91a961c11", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3242, + "public_pair_count": 2327, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3242, + "baseline_public_raw": 2327, + "candidate_precomputed": 3242, + "candidate_public_raw": 2327 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2230 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:adjacent_9ca0_d352_splitk_boundary_b1_n512_k8192_d352", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 8.04070981210856, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 3.774829931972789, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.7746562852159196, + "including_init_synchronized_e2e_speedup": 0.09279992951338945, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 3.0135299306029393, + "including_init_synchronized_e2e_speedup": 0.09970298608111366, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 3.5501464272362075, + "including_init_synchronized_e2e_speedup": 0.1673361493964697, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 3.746914382334343, + "including_init_synchronized_e2e_speedup": 0.7280359968707143, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 5.803412462908012, + "hot_synchronized_e2e_speedup": 3.774829931972789, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 935205, + "selected_route": "d352_exactd_splitk_c95c_v2", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_9ca0_d352_splitk_boundary_b1_n512_k8192_d352", + "source": "request_specific", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 416, + "K": 4096, + "N": 1024, + "baseline_07cf_adapter_bench_iters": 3098, + "baseline_07cf_adapter_gpu_span_ms": 0.15388849999999998, + "baseline_07cf_adapter_host_enqueue_ms": 0.15211200000000002, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.035264, + "baseline_07cf_adapter_kernel_sum_ms": 0.118592, + "baseline_07cf_adapter_submission_ms": 0.15211200000000002, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.256449, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.118592 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.235872, + "synchronized_e2e_ms": 0.331456 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.15388849999999998 + }, + "host_enqueue_ms": { + "median": 0.15211200000000002 + }, + "inter_kernel_gap_ms": { + "median": 0.035264 + }, + "kernel_sum_ms": { + "median": 0.118592 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3098, + "submission_ms": { + "median": 0.15211200000000002 + }, + "synchronized_e2e_ms": { + "median": 0.256449 + } + }, + "baseline_07cf_precomputed_bench_iters": 3167, + "baseline_07cf_precomputed_gpu_span_ms": 0.123008, + "baseline_07cf_precomputed_host_enqueue_ms": 0.045216, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.123008, + "baseline_07cf_precomputed_submission_ms": 0.045216, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.173761, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.123008 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.048768, + "synchronized_e2e_ms": 0.151808 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.123008 + }, + "host_enqueue_ms": { + "median": 0.045216 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.123008 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3167, + "submission_ms": { + "median": 0.045216 + }, + "synchronized_e2e_ms": { + "median": 0.173761 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.2510446475026014, + "submission": 3.36411889596603, + "synchronized_e2e": 1.4758720311232094 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.180135, + "after_init_synchronized_e2e_ms_per_call": 7.268103, + "including_init_host_enqueue_ms_per_call": 41.629579, + "including_init_synchronized_e2e_ms_per_call": 491.309597, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8549143000000001, + "after_init_synchronized_e2e_ms_per_call": 0.9576144, + "including_init_host_enqueue_ms_per_call": 4.2998587, + "including_init_synchronized_e2e_ms_per_call": 49.3617638, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22239223000000002, + "after_init_synchronized_e2e_ms_per_call": 0.32656554, + "including_init_host_enqueue_ms_per_call": 0.56688667, + "including_init_synchronized_e2e_ms_per_call": 5.166980479999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15914002300000005, + "after_init_synchronized_e2e_ms_per_call": 0.26346065399999996, + "including_init_host_enqueue_ms_per_call": 0.19358946700000004, + "including_init_synchronized_e2e_ms_per_call": 0.747502148, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.180135, + "synchronized_e2e_ms": 7.268103, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 43.792301, + "median": 0.15388849999999998, + "min": 0.145024, + "p90": 0.1700896, + "sample_count": 3098 + }, + "host_enqueue_ms": { + "max": 43.988909, + "median": 0.15211200000000002, + "min": 0.125793, + "p90": 0.1868672, + "sample_count": 3098 + }, + "sample_count": 3098, + "synchronized_e2e_ms": { + "max": 44.097261, + "median": 0.256449, + "min": 0.233537, + "p90": 0.2881696, + "sample_count": 3098 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 4, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3167, + "candidate_precomputed_gpu_span_ms": 0.028128, + "candidate_precomputed_host_enqueue_ms": 0.064448, + "candidate_precomputed_inter_kernel_gap_ms": 0.006208, + "candidate_precomputed_kernel_sum_ms": 0.021856, + "candidate_precomputed_launch_count": 3, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.064448, + "candidate_precomputed_synchronized_e2e_ms": 0.077344, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.021856 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.069792, + "synchronized_e2e_ms": 0.090304 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 3.0 + }, + "gpu_span_ms": { + "median": 0.028128 + }, + "host_enqueue_ms": { + "median": 0.064448 + }, + "inter_kernel_gap_ms": { + "median": 0.006208 + }, + "kernel_sum_ms": { + "median": 0.021856 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3167, + "submission_ms": { + "median": 0.064448 + }, + "synchronized_e2e_ms": { + "median": 0.077344 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295bef30", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc0404cd40" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.154721274175199, + "submission": 0.7591857000993047, + "synchronized_e2e": 1.0566818369880016 + }, + "candidate_public_raw_assignment_launch_count": 3, + "candidate_public_raw_bench_iters": 3098, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.03248, + "candidate_public_raw_host_enqueue_ms": 0.048928, + "candidate_public_raw_inter_kernel_gap_ms": 9.6e-05, + "candidate_public_raw_kernel_sum_ms": 0.032384, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.048928, + "candidate_public_raw_synchronized_e2e_ms": 0.081728, + "candidate_public_raw_tflops_from_gpu_span": 107.44029950738916, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.032384 + }, + "activity_count": { + "median": 4.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.062048, + "synchronized_e2e_ms": 0.0832 + }, + "correlated_kernel_activity_count": { + "median": 4.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.03248 + }, + "host_enqueue_ms": { + "median": 0.048928 + }, + "inter_kernel_gap_ms": { + "median": 9.6e-05 + }, + "kernel_sum_ms": { + "median": 0.032384 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3098, + "submission_ms": { + "median": 0.048928 + }, + "synchronized_e2e_ms": { + "median": 0.081728 + } + }, + "candidate_public_raw_total_launch_count": 4, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.333283, + "after_init_synchronized_e2e_ms_per_call": 2.363363, + "including_init_host_enqueue_ms_per_call": 37.207367, + "including_init_synchronized_e2e_ms_per_call": 37.306439, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2773635, + "after_init_synchronized_e2e_ms_per_call": 0.3098915, + "including_init_host_enqueue_ms_per_call": 3.7647718999999995, + "including_init_synchronized_e2e_ms_per_call": 3.8041990999999995, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07177155, + "after_init_synchronized_e2e_ms_per_call": 0.10454434999999998, + "including_init_host_enqueue_ms_per_call": 0.42051238999999996, + "including_init_synchronized_e2e_ms_per_call": 0.45397510999999996, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.051212355, + "after_init_synchronized_e2e_ms_per_call": 0.084009635, + "including_init_host_enqueue_ms_per_call": 0.086086439, + "including_init_synchronized_e2e_ms_per_call": 0.11895271099999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.333283, + "synchronized_e2e_ms": 2.363363, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.033344, + "median": 0.03248, + "min": 0.031904, + "p90": 0.032768, + "sample_count": 3098 + }, + "host_enqueue_ms": { + "max": 0.471937, + "median": 0.048928, + "min": 0.03888, + "p90": 0.06701760000000001, + "sample_count": 3098 + }, + "sample_count": 3098, + "synchronized_e2e_ms": { + "max": 0.530881, + "median": 0.081728, + "min": 0.07264, + "p90": 0.09819520000000001, + "sample_count": 3098 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.071648, + "submission_ms": 0.071648, + "synchronized_e2e_ms": 0.173984 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.048768, + "submission_ms": 0.048768, + "synchronized_e2e_ms": 0.151808 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.235872, + "submission_ms": 0.235872, + "synchronized_e2e_ms": 0.331456 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.180135, + "submission_ms": 7.180135, + "synchronized_e2e_ms": 7.268103 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.07872, + "submission_ms": 0.07872, + "synchronized_e2e_ms": 0.098016 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.069792, + "submission_ms": 0.069792, + "synchronized_e2e_ms": 0.090304 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.281089, + "submission_ms": 1.281089, + "synchronized_e2e_ms": 1.309857 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.193346, + "submission_ms": 1.193346, + "synchronized_e2e_ms": 1.219362 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.062048, + "submission_ms": 0.062048, + "synchronized_e2e_ms": 0.0832 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.333283, + "submission_ms": 2.333283, + "synchronized_e2e_ms": 2.363363 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 41.154087, + "evolution_kernel_ms": 0.274271, + "evolution_speedup": 150.049, + "evolution_tflops": 12.7234, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_9ca0_d416_splitk_min_b1_n1024_k4096_d416", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 12530, + "measurement_schedule_sha256": "7118dca39cce6adcb0e1a236501fb681dcec5e3109c51860ee4969ba9ff44e43", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3167, + "public_pair_count": 3098, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3167, + "baseline_public_raw": 3098, + "candidate_precomputed": 3167, + "candidate_public_raw": 3098 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2508 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:adjacent_9ca0_d416_splitk_min_b1_n1024_k4096_d416", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 4.373151308304892, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 3.1378352584181677, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.0753223267013996, + "including_init_synchronized_e2e_speedup": 13.16956563450079, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 3.090160265770439, + "including_init_synchronized_e2e_speedup": 12.975599463235246, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 3.123703385214027, + "including_init_synchronized_e2e_speedup": 11.3816382576569, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 3.1360766416852064, + "including_init_synchronized_e2e_speedup": 6.284027843636116, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 4.737946428571428, + "hot_synchronized_e2e_speedup": 3.1378352584181677, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 941605, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_9ca0_d416_splitk_min_b1_n1024_k4096_d416", + "source": "guard_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 480, + "K": 512, + "N": 4096, + "baseline_07cf_adapter_bench_iters": 2508, + "baseline_07cf_adapter_gpu_span_ms": 0.064272, + "baseline_07cf_adapter_host_enqueue_ms": 0.159424, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.029568, + "baseline_07cf_adapter_kernel_sum_ms": 0.03472, + "baseline_07cf_adapter_submission_ms": 0.159424, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.18004799999999999, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.03472 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.217696, + "synchronized_e2e_ms": 0.24 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.064272 + }, + "host_enqueue_ms": { + "median": 0.159424 + }, + "inter_kernel_gap_ms": { + "median": 0.029568 + }, + "kernel_sum_ms": { + "median": 0.03472 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2508, + "submission_ms": { + "median": 0.159424 + }, + "synchronized_e2e_ms": { + "median": 0.18004799999999999 + } + }, + "baseline_07cf_precomputed_bench_iters": 4452, + "baseline_07cf_precomputed_gpu_span_ms": 0.022528, + "baseline_07cf_precomputed_host_enqueue_ms": 0.044192, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.022528, + "baseline_07cf_precomputed_submission_ms": 0.044192, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.072224, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.022528 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.052288, + "synchronized_e2e_ms": 0.07312 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.022528 + }, + "host_enqueue_ms": { + "median": 0.044192 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.022528 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4452, + "submission_ms": { + "median": 0.044192 + }, + "synchronized_e2e_ms": { + "median": 0.072224 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.8529829545454546, + "submission": 3.607530774800869, + "synchronized_e2e": 2.4929109437306156 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.848008, + "after_init_synchronized_e2e_ms_per_call": 7.879272, + "including_init_host_enqueue_ms_per_call": 43.399276, + "including_init_synchronized_e2e_ms_per_call": 43.510509, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9282824, + "after_init_synchronized_e2e_ms_per_call": 0.9499704, + "including_init_host_enqueue_ms_per_call": 4.4834092, + "including_init_synchronized_e2e_ms_per_call": 4.5130941, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23630984000000002, + "after_init_synchronized_e2e_ms_per_call": 0.25704024000000003, + "including_init_host_enqueue_ms_per_call": 0.59182252, + "including_init_synchronized_e2e_ms_per_call": 0.6133526100000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.167112584, + "after_init_synchronized_e2e_ms_per_call": 0.18774722399999996, + "including_init_host_enqueue_ms_per_call": 0.20266385200000003, + "including_init_synchronized_e2e_ms_per_call": 0.22337846099999997, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.848008, + "synchronized_e2e_ms": 7.879272, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.240645, + "median": 0.064272, + "min": 0.054176, + "p90": 0.09256960000000003, + "sample_count": 2508 + }, + "host_enqueue_ms": { + "max": 34.697252, + "median": 0.159424, + "min": 0.124608, + "p90": 0.27041919999999997, + "sample_count": 2508 + }, + "sample_count": 2508, + "synchronized_e2e_ms": { + "max": 35.566725, + "median": 0.18004799999999999, + "min": 0.141568, + "p90": 0.3014535000000006, + "sample_count": 2508 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 4452, + "candidate_precomputed_gpu_span_ms": 0.0264, + "candidate_precomputed_host_enqueue_ms": 0.05392, + "candidate_precomputed_inter_kernel_gap_ms": 0.002432, + "candidate_precomputed_kernel_sum_ms": 0.023936, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.05392, + "candidate_precomputed_synchronized_e2e_ms": 0.066112, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.023936 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.043104, + "synchronized_e2e_ms": 0.060288 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.0264 + }, + "host_enqueue_ms": { + "median": 0.05392 + }, + "inter_kernel_gap_ms": { + "median": 0.002432 + }, + "kernel_sum_ms": { + "median": 0.023936 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4452, + "submission_ms": { + "median": 0.05392 + }, + "synchronized_e2e_ms": { + "median": 0.066112 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01230680", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc012328d0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.5127272727272727, + "submission": 0.8712166172106824, + "synchronized_e2e": 1.3272023233301065 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2508, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.039936, + "candidate_public_raw_host_enqueue_ms": 0.046976, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.04016, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.046976, + "candidate_public_raw_synchronized_e2e_ms": 0.087744, + "candidate_public_raw_tflops_from_gpu_span": 100.82461538461538, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.039936 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.082688, + "synchronized_e2e_ms": 0.104544 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.039936 + }, + "host_enqueue_ms": { + "median": 0.046976 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.04016 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2508, + "submission_ms": { + "median": 0.046976 + }, + "synchronized_e2e_ms": { + "median": 0.087744 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.276835, + "after_init_synchronized_e2e_ms_per_call": 2.300419, + "including_init_host_enqueue_ms_per_call": 38.114855, + "including_init_synchronized_e2e_ms_per_call": 486.400632, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.26996190000000003, + "after_init_synchronized_e2e_ms_per_call": 0.30901150000000005, + "including_init_host_enqueue_ms_per_call": 3.8537638999999997, + "including_init_synchronized_e2e_ms_per_call": 48.719032799999994, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06927459, + "after_init_synchronized_e2e_ms_per_call": 0.10987075, + "including_init_host_enqueue_ms_per_call": 0.42765479, + "including_init_synchronized_e2e_ms_per_call": 4.95087288, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.049205859, + "after_init_synchronized_e2e_ms_per_call": 0.089956675, + "including_init_host_enqueue_ms_per_call": 0.085043879, + "including_init_synchronized_e2e_ms_per_call": 0.574056888, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.276835, + "synchronized_e2e_ms": 2.300419, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.040673, + "median": 0.039936, + "min": 0.039552, + "p90": 0.04016, + "sample_count": 2508 + }, + "host_enqueue_ms": { + "max": 48.600498, + "median": 0.046976, + "min": 0.035616, + "p90": 0.08609920000000001, + "sample_count": 2508 + }, + "sample_count": 2508, + "synchronized_e2e_ms": { + "max": 48.751251, + "median": 0.087744, + "min": 0.078048, + "p90": 0.11937030000000001, + "sample_count": 2508 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.084321, + "submission_ms": 0.084321, + "synchronized_e2e_ms": 0.106177 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.052288, + "submission_ms": 0.052288, + "synchronized_e2e_ms": 0.07312 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.217696, + "submission_ms": 0.217696, + "synchronized_e2e_ms": 0.24 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.848008, + "submission_ms": 7.848008, + "synchronized_e2e_ms": 7.879272 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.082944, + "submission_ms": 0.082944, + "synchronized_e2e_ms": 0.102944 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.043104, + "submission_ms": 0.043104, + "synchronized_e2e_ms": 0.060288 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.418114, + "submission_ms": 1.418114, + "synchronized_e2e_ms": 1.45101 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.246017, + "submission_ms": 1.246017, + "synchronized_e2e_ms": 1.277505 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.082688, + "submission_ms": 0.082688, + "synchronized_e2e_ms": 0.104544 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.276835, + "submission_ms": 2.276835, + "synchronized_e2e_ms": 2.300419 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.447807, + "evolution_kernel_ms": 0.185087, + "evolution_speedup": 2.4194, + "evolution_tflops": 21.7548, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_9ca0_d480_splitd_large_n_b2_n4096_k512_d480", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 13920, + "measurement_schedule_sha256": "d5a4fd6d20c67a366a18905bb066b867de519e903acec9469c9a8461772ae7f4", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 4452, + "public_pair_count": 2508, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 4452, + "baseline_public_raw": 2508, + "candidate_precomputed": 4452, + "candidate_public_raw": 2508 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2786 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:adjacent_9ca0_d480_splitd_large_n_b2_n4096_k512_d480", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.8533333333333333, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.0519693654266957, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.4251464624487973, + "including_init_synchronized_e2e_speedup": 0.08945405523239534, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 3.074223451230779, + "including_init_synchronized_e2e_speedup": 0.09263513334772115, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.339478341596831, + "including_init_synchronized_e2e_speedup": 0.1238877718872071, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.0870849661795523, + "including_init_synchronized_e2e_speedup": 0.3891225167217225, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.609375, + "hot_synchronized_e2e_speedup": 2.0519693654266957, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 948005, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_9ca0_d480_splitd_large_n_b2_n4096_k512_d480", + "source": "random_legal", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 6, + "D": 48, + "K": 512, + "N": 12288, + "baseline_07cf_adapter_bench_iters": 1518, + "baseline_07cf_adapter_gpu_span_ms": 0.07728, + "baseline_07cf_adapter_host_enqueue_ms": 0.173184, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.008992, + "baseline_07cf_adapter_kernel_sum_ms": 0.068224, + "baseline_07cf_adapter_submission_ms": 0.173184, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.1952325, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.068224 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.21136, + "synchronized_e2e_ms": 0.23152 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.07728 + }, + "host_enqueue_ms": { + "median": 0.173184 + }, + "inter_kernel_gap_ms": { + "median": 0.008992 + }, + "kernel_sum_ms": { + "median": 0.068224 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1518, + "submission_ms": { + "median": 0.173184 + }, + "synchronized_e2e_ms": { + "median": 0.1952325 + } + }, + "baseline_07cf_precomputed_bench_iters": 3730, + "baseline_07cf_precomputed_gpu_span_ms": 0.026688, + "baseline_07cf_precomputed_host_enqueue_ms": 0.047264, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.026688, + "baseline_07cf_precomputed_submission_ms": 0.047264, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.079136, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.026688 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.043296, + "synchronized_e2e_ms": 0.074368 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.026688 + }, + "host_enqueue_ms": { + "median": 0.047264 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.026688 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3730, + "submission_ms": { + "median": 0.047264 + }, + "synchronized_e2e_ms": { + "median": 0.079136 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.8956834532374103, + "submission": 3.6641841570751525, + "synchronized_e2e": 2.4670503942579862 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.621289, + "after_init_synchronized_e2e_ms_per_call": 8.650921, + "including_init_host_enqueue_ms_per_call": 45.555727000000005, + "including_init_synchronized_e2e_ms_per_call": 463.76224, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 1.0179945, + "after_init_synchronized_e2e_ms_per_call": 1.0408013500000002, + "including_init_host_enqueue_ms_per_call": 4.7114383, + "including_init_synchronized_e2e_ms_per_call": 46.551933250000005, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.25766505, + "after_init_synchronized_e2e_ms_per_call": 0.279789385, + "including_init_host_enqueue_ms_per_call": 0.6270094300000001, + "including_init_synchronized_e2e_ms_per_call": 4.830902575, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.181632105, + "after_init_synchronized_e2e_ms_per_call": 0.20368818850000003, + "including_init_host_enqueue_ms_per_call": 0.21856654300000003, + "including_init_synchronized_e2e_ms_per_call": 0.6587995075, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.621289, + "synchronized_e2e_ms": 8.650921, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 6.576198, + "median": 0.07728, + "min": 0.068928, + "p90": 0.09614399999999998, + "sample_count": 1518 + }, + "host_enqueue_ms": { + "max": 17.968499, + "median": 0.173184, + "min": 0.140864, + "p90": 0.2165856, + "sample_count": 1518 + }, + "sample_count": 1518, + "synchronized_e2e_ms": { + "max": 23.1446, + "median": 0.1952325, + "min": 0.165344, + "p90": 0.2421184, + "sample_count": 1518 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3730, + "candidate_precomputed_gpu_span_ms": 0.028736, + "candidate_precomputed_host_enqueue_ms": 0.0608, + "candidate_precomputed_inter_kernel_gap_ms": 0.003456, + "candidate_precomputed_kernel_sum_ms": 0.02528, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.0608, + "candidate_precomputed_synchronized_e2e_ms": 0.07384, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.02528 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.067904, + "synchronized_e2e_ms": 0.085152 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.028736 + }, + "host_enqueue_ms": { + "median": 0.0608 + }, + "inter_kernel_gap_ms": { + "median": 0.003456 + }, + "kernel_sum_ms": { + "median": 0.02528 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3730, + "submission_ms": { + "median": 0.0608 + }, + "synchronized_e2e_ms": { + "median": 0.07384 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc028aee40", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc028aef00" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.3062360801781736, + "submission": 0.871578947368421, + "synchronized_e2e": 1.6138678223185263 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 1518, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.066272, + "candidate_public_raw_host_enqueue_ms": 0.052992, + "candidate_public_raw_inter_kernel_gap_ms": 0.000128, + "candidate_public_raw_kernel_sum_ms": 0.066144, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.052992, + "candidate_public_raw_synchronized_e2e_ms": 0.119168, + "candidate_public_raw_tflops_from_gpu_span": 54.6818966682762, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.066144 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.105825, + "synchronized_e2e_ms": 0.140673 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.066272 + }, + "host_enqueue_ms": { + "median": 0.052992 + }, + "inter_kernel_gap_ms": { + "median": 0.000128 + }, + "kernel_sum_ms": { + "median": 0.066144 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1518, + "submission_ms": { + "median": 0.052992 + }, + "synchronized_e2e_ms": { + "median": 0.119168 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.579298, + "after_init_synchronized_e2e_ms_per_call": 2.622722, + "including_init_host_enqueue_ms_per_call": 39.937321, + "including_init_synchronized_e2e_ms_per_call": 40.058345, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.3056226, + "after_init_synchronized_e2e_ms_per_call": 0.3695234, + "including_init_host_enqueue_ms_per_call": 4.0414249, + "including_init_synchronized_e2e_ms_per_call": 4.113085700000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07825505999999999, + "after_init_synchronized_e2e_ms_per_call": 0.14420354, + "including_init_host_enqueue_ms_per_call": 0.45183528999999995, + "including_init_synchronized_e2e_ms_per_call": 0.51855977, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.055518305999999996, + "after_init_synchronized_e2e_ms_per_call": 0.12167155399999999, + "including_init_host_enqueue_ms_per_call": 0.092876329, + "including_init_synchronized_e2e_ms_per_call": 0.159107177, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.579298, + "synchronized_e2e_ms": 2.622722, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.067168, + "median": 0.066272, + "min": 0.065792, + "p90": 0.066464, + "sample_count": 1518 + }, + "host_enqueue_ms": { + "max": 2.114594, + "median": 0.052992, + "min": 0.041632, + "p90": 0.0721923, + "sample_count": 1518 + }, + "sample_count": 1518, + "synchronized_e2e_ms": { + "max": 2.352963, + "median": 0.119168, + "min": 0.109984, + "p90": 0.1357952, + "sample_count": 1518 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.071776, + "submission_ms": 0.071776, + "synchronized_e2e_ms": 0.100768 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.043296, + "submission_ms": 0.043296, + "synchronized_e2e_ms": 0.074368 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.21136, + "submission_ms": 0.21136, + "synchronized_e2e_ms": 0.23152 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.621289, + "submission_ms": 8.621289, + "synchronized_e2e_ms": 8.650921 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.079104, + "submission_ms": 0.079104, + "synchronized_e2e_ms": 0.097376 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.067904, + "submission_ms": 0.067904, + "synchronized_e2e_ms": 0.085152 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.336161, + "submission_ms": 1.336161, + "synchronized_e2e_ms": 1.362241 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.183969, + "submission_ms": 1.183969, + "synchronized_e2e_ms": 1.209793 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.105825, + "submission_ms": 0.105825, + "synchronized_e2e_ms": 0.140673 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.579298, + "submission_ms": 2.579298, + "synchronized_e2e_ms": 2.622722 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.318304, + "evolution_kernel_ms": 0.182336, + "evolution_speedup": 1.7457, + "evolution_tflops": 19.8747, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_9ca0_d48_guard_boundary_b6_n12288_k512_d48", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 10496, + "measurement_schedule_sha256": "740a7a64892cc3af481df1f543a8265567b7ad2347a147f01e48e7ae8f405d2d", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3730, + "public_pair_count": 1518, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3730, + "baseline_public_raw": 1518, + "candidate_precomputed": 3730, + "candidate_public_raw": 1518 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2100 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:adjacent_9ca0_d48_guard_boundary_b6_n12288_k512_d48", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.9287305122494431, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.6382963547261011, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.298451379902254, + "including_init_synchronized_e2e_speedup": 11.57716925150053, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.8166047129897596, + "including_init_synchronized_e2e_speedup": 11.318007122973391, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 1.9402393658297155, + "including_init_synchronized_e2e_speedup": 9.315999532705748, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.6740822468660181, + "including_init_synchronized_e2e_speedup": 4.1406020766743925, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.166103331723805, + "hot_synchronized_e2e_speedup": 1.6382963547261011, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 924805, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_9ca0_d48_guard_boundary_b6_n12288_k512_d48", + "source": "guard_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 64, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 112, + "K": 256, + "N": 3456, + "baseline_07cf_adapter_bench_iters": 4899, + "baseline_07cf_adapter_gpu_span_ms": 0.0536, + "baseline_07cf_adapter_host_enqueue_ms": 0.152448, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.03472, + "baseline_07cf_adapter_kernel_sum_ms": 0.01888, + "baseline_07cf_adapter_submission_ms": 0.152448, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.172609, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.01888 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.196417, + "synchronized_e2e_ms": 0.214465 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.0536 + }, + "host_enqueue_ms": { + "median": 0.152448 + }, + "inter_kernel_gap_ms": { + "median": 0.03472 + }, + "kernel_sum_ms": { + "median": 0.01888 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 4899, + "submission_ms": { + "median": 0.152448 + }, + "synchronized_e2e_ms": { + "median": 0.172609 + } + }, + "baseline_07cf_precomputed_bench_iters": 11262, + "baseline_07cf_precomputed_gpu_span_ms": 0.008928, + "baseline_07cf_precomputed_host_enqueue_ms": 0.04256, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.008928, + "baseline_07cf_precomputed_submission_ms": 0.04256, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.058176, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.008928 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.046049, + "synchronized_e2e_ms": 0.061185 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.008928 + }, + "host_enqueue_ms": { + "median": 0.04256 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.008928 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 11262, + "submission_ms": { + "median": 0.04256 + }, + "synchronized_e2e_ms": { + "median": 0.058176 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 6.003584229390681, + "submission": 3.581954887218045, + "synchronized_e2e": 2.9670138888888893 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.601192, + "after_init_synchronized_e2e_ms_per_call": 7.624968, + "including_init_host_enqueue_ms_per_call": 41.829099, + "including_init_synchronized_e2e_ms_per_call": 41.970411000000006, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8973224, + "after_init_synchronized_e2e_ms_per_call": 0.9178449000000001, + "including_init_host_enqueue_ms_per_call": 4.3201130999999995, + "including_init_synchronized_e2e_ms_per_call": 4.3523892, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22693544, + "after_init_synchronized_e2e_ms_per_call": 0.24713259, + "including_init_host_enqueue_ms_per_call": 0.56921451, + "including_init_synchronized_e2e_ms_per_call": 0.5905870200000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15989674399999998, + "after_init_synchronized_e2e_ms_per_call": 0.180061359, + "including_init_host_enqueue_ms_per_call": 0.19412465099999998, + "including_init_synchronized_e2e_ms_per_call": 0.21440680200000004, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.601192, + "synchronized_e2e_ms": 7.624968, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.628742, + "median": 0.0536, + "min": 0.043168, + "p90": 0.0862848, + "sample_count": 4899 + }, + "host_enqueue_ms": { + "max": 54.077688, + "median": 0.152448, + "min": 0.122816, + "p90": 0.26858879999999996, + "sample_count": 4899 + }, + "sample_count": 4899, + "synchronized_e2e_ms": { + "max": 54.236152, + "median": 0.172609, + "min": 0.139488, + "p90": 0.2982786, + "sample_count": 4899 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 11262, + "candidate_precomputed_gpu_span_ms": 0.013792, + "candidate_precomputed_host_enqueue_ms": 0.05376, + "candidate_precomputed_inter_kernel_gap_ms": 0.004192, + "candidate_precomputed_kernel_sum_ms": 0.009696, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.05376, + "candidate_precomputed_synchronized_e2e_ms": 0.066272, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.009696 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.04624, + "synchronized_e2e_ms": 0.061376 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.013792 + }, + "host_enqueue_ms": { + "median": 0.05376 + }, + "inter_kernel_gap_ms": { + "median": 0.004192 + }, + "kernel_sum_ms": { + "median": 0.009696 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 11262, + "submission_ms": { + "median": 0.05376 + }, + "synchronized_e2e_ms": { + "median": 0.066272 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffe4d6f5f70", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc03924bc0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.5011600928074245, + "submission": 0.843452380952381, + "synchronized_e2e": 1.0304200869145341 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 4899, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.020704, + "candidate_public_raw_host_enqueue_ms": 0.045344, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.020896, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.045344, + "candidate_public_raw_synchronized_e2e_ms": 0.068288, + "candidate_public_raw_tflops_from_gpu_span": 38.288420401854715, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.020704 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.069216, + "synchronized_e2e_ms": 0.088576 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.020704 + }, + "host_enqueue_ms": { + "median": 0.045344 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.020896 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 4899, + "submission_ms": { + "median": 0.045344 + }, + "synchronized_e2e_ms": { + "median": 0.068288 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.486978, + "after_init_synchronized_e2e_ms_per_call": 2.513154, + "including_init_host_enqueue_ms_per_call": 37.064518, + "including_init_synchronized_e2e_ms_per_call": 450.60094599999996, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2895074, + "after_init_synchronized_e2e_ms_per_call": 0.3127746, + "including_init_host_enqueue_ms_per_call": 3.7472614, + "including_init_synchronized_e2e_ms_per_call": 45.1215538, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06976034, + "after_init_synchronized_e2e_ms_per_call": 0.09273666, + "including_init_host_enqueue_ms_per_call": 0.41553574, + "including_init_synchronized_e2e_ms_per_call": 4.57361458, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.047785634, + "after_init_synchronized_e2e_ms_per_call": 0.070732866, + "including_init_host_enqueue_ms_per_call": 0.082363174, + "including_init_synchronized_e2e_ms_per_call": 0.518820658, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.486978, + "synchronized_e2e_ms": 2.513154, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.021344, + "median": 0.020704, + "min": 0.02, + "p90": 0.020992, + "sample_count": 4899 + }, + "host_enqueue_ms": { + "max": 0.641984, + "median": 0.045344, + "min": 0.034688, + "p90": 0.08199039999999999, + "sample_count": 4899 + }, + "sample_count": 4899, + "synchronized_e2e_ms": { + "max": 0.802049, + "median": 0.068288, + "min": 0.059264, + "p90": 0.1069504, + "sample_count": 4899 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.063552, + "submission_ms": 0.063552, + "synchronized_e2e_ms": 0.079648 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.046049, + "submission_ms": 0.046049, + "synchronized_e2e_ms": 0.061185 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.196417, + "submission_ms": 0.196417, + "synchronized_e2e_ms": 0.214465 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.601192, + "submission_ms": 7.601192, + "synchronized_e2e_ms": 7.624968 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.071552, + "submission_ms": 0.071552, + "synchronized_e2e_ms": 0.088256 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.04624, + "submission_ms": 0.04624, + "synchronized_e2e_ms": 0.061376 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.273313, + "submission_ms": 1.273313, + "synchronized_e2e_ms": 1.298241 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.111745, + "submission_ms": 1.111745, + "synchronized_e2e_ms": 1.133473 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.069216, + "submission_ms": 0.069216, + "synchronized_e2e_ms": 0.088576 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.486978, + "submission_ms": 2.486978, + "synchronized_e2e_ms": 2.513154 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.231359, + "evolution_kernel_ms": 0.168912, + "evolution_speedup": 1.3697, + "evolution_tflops": 4.6931, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_a2f8_d112_lowk_tail_b4_n3456_k256_d112", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 32322, + "measurement_schedule_sha256": "dbca175a7e4213eed7d68347f3a3d4f04a8a7c3adb2fc3dac6b957da87146400", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 11262, + "public_pair_count": 4899, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 11262, + "baseline_public_raw": 4899, + "candidate_precomputed": 11262, + "candidate_public_raw": 4899 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 6466 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:adjacent_a2f8_d112_lowk_tail_b4_n3456_k256_d112", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.6473317865429234, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.5276622539831304, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.0340233825702683, + "including_init_synchronized_e2e_speedup": 0.09314319326795202, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.934525054144422, + "including_init_synchronized_e2e_speedup": 0.09645920482463527, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.6648856018752456, + "including_init_synchronized_e2e_speedup": 0.12912916243152261, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.545653374203726, + "including_init_synchronized_e2e_speedup": 0.4132580279792946, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.58887171561051, + "hot_synchronized_e2e_speedup": 2.5276622539831304, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 1028112, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_a2f8_d112_lowk_tail_b4_n3456_k256_d112", + "source": "tail_divisibility", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 6, + "D": 128, + "K": 512, + "N": 8576, + "baseline_07cf_adapter_bench_iters": 2011, + "baseline_07cf_adapter_gpu_span_ms": 0.0744, + "baseline_07cf_adapter_host_enqueue_ms": 0.161376, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.017184, + "baseline_07cf_adapter_kernel_sum_ms": 0.057248, + "baseline_07cf_adapter_submission_ms": 0.161376, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.183744, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.057248 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.246721, + "synchronized_e2e_ms": 0.268289 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.0744 + }, + "host_enqueue_ms": { + "median": 0.161376 + }, + "inter_kernel_gap_ms": { + "median": 0.017184 + }, + "kernel_sum_ms": { + "median": 0.057248 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2011, + "submission_ms": { + "median": 0.161376 + }, + "synchronized_e2e_ms": { + "median": 0.183744 + } + }, + "baseline_07cf_precomputed_bench_iters": 5356, + "baseline_07cf_precomputed_gpu_span_ms": 0.026432, + "baseline_07cf_precomputed_host_enqueue_ms": 0.045696, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.026432, + "baseline_07cf_precomputed_submission_ms": 0.045696, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.077536, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.026432 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.053184, + "synchronized_e2e_ms": 0.082784 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.026432 + }, + "host_enqueue_ms": { + "median": 0.045696 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.026432 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5356, + "submission_ms": { + "median": 0.045696 + }, + "synchronized_e2e_ms": { + "median": 0.077536 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.8147699757869247, + "submission": 3.5315126050420167, + "synchronized_e2e": 2.369789517127528 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.400168, + "after_init_synchronized_e2e_ms_per_call": 7.428808, + "including_init_host_enqueue_ms_per_call": 41.849612, + "including_init_synchronized_e2e_ms_per_call": 491.470302, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8852551999999999, + "after_init_synchronized_e2e_ms_per_call": 0.9082504, + "including_init_host_enqueue_ms_per_call": 4.3301996, + "including_init_synchronized_e2e_ms_per_call": 49.3123998, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23376392, + "after_init_synchronized_e2e_ms_per_call": 0.25619464000000003, + "including_init_host_enqueue_ms_per_call": 0.57825836, + "including_init_synchronized_e2e_ms_per_call": 5.09660958, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.16861479199999999, + "after_init_synchronized_e2e_ms_per_call": 0.190989064, + "including_init_host_enqueue_ms_per_call": 0.20306423599999998, + "including_init_synchronized_e2e_ms_per_call": 0.6750305579999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.400168, + "synchronized_e2e_ms": 7.428808, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.362432, + "median": 0.0744, + "min": 0.063168, + "p90": 0.095296, + "sample_count": 2011 + }, + "host_enqueue_ms": { + "max": 27.455548, + "median": 0.161376, + "min": 0.132064, + "p90": 0.203872, + "sample_count": 2011 + }, + "sample_count": 2011, + "synchronized_e2e_ms": { + "max": 27.497116, + "median": 0.183744, + "min": 0.154784, + "p90": 0.229441, + "sample_count": 2011 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 5356, + "candidate_precomputed_gpu_span_ms": 0.018816, + "candidate_precomputed_host_enqueue_ms": 0.04192, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.018816, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.04192, + "candidate_precomputed_synchronized_e2e_ms": 0.054432, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.018816 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.038112, + "synchronized_e2e_ms": 0.054176 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.018816 + }, + "host_enqueue_ms": { + "median": 0.04192 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.018816 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5356, + "submission_ms": { + "median": 0.04192 + }, + "synchronized_e2e_ms": { + "median": 0.054432 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282c65a0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04dcba10" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.6462585034013606, + "submission": 1.251908396946565, + "synchronized_e2e": 1.8741916519694297 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 2011, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.049792, + "candidate_public_raw_host_enqueue_ms": 0.05248, + "candidate_public_raw_inter_kernel_gap_ms": 0.00016, + "candidate_public_raw_kernel_sum_ms": 0.049632, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.05248, + "candidate_public_raw_synchronized_e2e_ms": 0.102016, + "candidate_public_raw_tflops_from_gpu_span": 135.45229820051412, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.049632 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.276672, + "synchronized_e2e_ms": 0.313088 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.049792 + }, + "host_enqueue_ms": { + "median": 0.05248 + }, + "inter_kernel_gap_ms": { + "median": 0.00016 + }, + "kernel_sum_ms": { + "median": 0.049632 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2011, + "submission_ms": { + "median": 0.05248 + }, + "synchronized_e2e_ms": { + "median": 0.102016 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.370722, + "after_init_synchronized_e2e_ms_per_call": 2.399298, + "including_init_host_enqueue_ms_per_call": 37.244806, + "including_init_synchronized_e2e_ms_per_call": 37.342374, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2843042, + "after_init_synchronized_e2e_ms_per_call": 0.3317442, + "including_init_host_enqueue_ms_per_call": 3.7717126, + "including_init_synchronized_e2e_ms_per_call": 3.8260517999999997, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07566242000000001, + "after_init_synchronized_e2e_ms_per_call": 0.12498882, + "including_init_host_enqueue_ms_per_call": 0.42440326, + "including_init_synchronized_e2e_ms_per_call": 0.47441958, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.054798242000000004, + "after_init_synchronized_e2e_ms_per_call": 0.10431328200000001, + "including_init_host_enqueue_ms_per_call": 0.089672326, + "including_init_synchronized_e2e_ms_per_call": 0.139256358, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.370722, + "synchronized_e2e_ms": 2.399298, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.050656, + "median": 0.049792, + "min": 0.049088, + "p90": 0.050112, + "sample_count": 2011 + }, + "host_enqueue_ms": { + "max": 45.973359, + "median": 0.05248, + "min": 0.038912, + "p90": 0.071872, + "sample_count": 2011 + }, + "sample_count": 2011, + "synchronized_e2e_ms": { + "max": 46.079727, + "median": 0.102016, + "min": 0.088928, + "p90": 0.119456, + "sample_count": 2011 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.086912, + "submission_ms": 0.086912, + "synchronized_e2e_ms": 0.11488 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.053184, + "submission_ms": 0.053184, + "synchronized_e2e_ms": 0.082784 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.246721, + "submission_ms": 0.246721, + "synchronized_e2e_ms": 0.268289 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.400168, + "submission_ms": 7.400168, + "synchronized_e2e_ms": 7.428808 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.059648, + "submission_ms": 0.059648, + "synchronized_e2e_ms": 0.078784 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.038112, + "submission_ms": 0.038112, + "synchronized_e2e_ms": 0.054176 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.264513, + "submission_ms": 1.264513, + "synchronized_e2e_ms": 1.287489 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.100833, + "submission_ms": 1.100833, + "synchronized_e2e_ms": 1.122113 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.276672, + "submission_ms": 0.276672, + "synchronized_e2e_ms": 0.313088 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.370722, + "submission_ms": 2.370722, + "synchronized_e2e_ms": 2.399298 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.349887, + "evolution_kernel_ms": 0.158688, + "evolution_speedup": 2.2049, + "evolution_tflops": 42.5013, + "expected_route": "aligned_weave_v10_fallback", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_a2f8_d128_odd_fallback_b6_n8576_k512_d128", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 14734, + "measurement_schedule_sha256": "4f87946802236cbe25cd77a4a9cebe50fb570bd939e10094a917fd8d7469c61c", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 5356, + "public_pair_count": 2011, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 5356, + "baseline_public_raw": 2011, + "candidate_precomputed": 5356, + "candidate_public_raw": 2011 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2950 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:adjacent_a2f8_d128_odd_fallback_b6_n8576_k512_d128", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.4047619047619049, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.801129234629862, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.0962423175445486, + "including_init_synchronized_e2e_speedup": 13.161195964670055, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.7378034039479817, + "including_init_synchronized_e2e_speedup": 12.88858655808058, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.0497404487857396, + "including_init_synchronized_e2e_speedup": 10.742831440473008, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.8309179841546925, + "including_init_synchronized_e2e_speedup": 4.847394888784898, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.4942159383033418, + "hot_synchronized_e2e_speedup": 1.801129234629862, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 1028128, + "selected_route": "aligned_weave_v10_fallback", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_a2f8_d128_odd_fallback_b6_n8576_k512_d128", + "source": "guard_miss_fallback", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 224, + "K": 768, + "N": 6144, + "baseline_07cf_adapter_bench_iters": 3039, + "baseline_07cf_adapter_gpu_span_ms": 0.06752, + "baseline_07cf_adapter_host_enqueue_ms": 0.149888, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.025984, + "baseline_07cf_adapter_kernel_sum_ms": 0.041568, + "baseline_07cf_adapter_submission_ms": 0.149888, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.170305, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.041568 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.21456, + "synchronized_e2e_ms": 0.232416 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.06752 + }, + "host_enqueue_ms": { + "median": 0.149888 + }, + "inter_kernel_gap_ms": { + "median": 0.025984 + }, + "kernel_sum_ms": { + "median": 0.041568 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3039, + "submission_ms": { + "median": 0.149888 + }, + "synchronized_e2e_ms": { + "median": 0.170305 + } + }, + "baseline_07cf_precomputed_bench_iters": 4838, + "baseline_07cf_precomputed_gpu_span_ms": 0.02528, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042272, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.02528, + "baseline_07cf_precomputed_submission_ms": 0.042272, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.073024, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.02528 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.047264, + "synchronized_e2e_ms": 0.075776 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.02528 + }, + "host_enqueue_ms": { + "median": 0.042272 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.02528 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4838, + "submission_ms": { + "median": 0.042272 + }, + "synchronized_e2e_ms": { + "median": 0.073024 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.670886075949367, + "submission": 3.545798637395912, + "synchronized_e2e": 2.332178461875548 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.417928, + "after_init_synchronized_e2e_ms_per_call": 7.444712, + "including_init_host_enqueue_ms_per_call": 42.969196, + "including_init_synchronized_e2e_ms_per_call": 43.075949, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8766919999999999, + "after_init_synchronized_e2e_ms_per_call": 0.8977457, + "including_init_host_enqueue_ms_per_call": 4.4318188, + "including_init_synchronized_e2e_ms_per_call": 4.4608694, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22256839999999997, + "after_init_synchronized_e2e_ms_per_call": 0.24304907, + "including_init_host_enqueue_ms_per_call": 0.57808108, + "including_init_synchronized_e2e_ms_per_call": 0.59936144, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15715604, + "after_init_synchronized_e2e_ms_per_call": 0.17757940700000002, + "including_init_host_enqueue_ms_per_call": 0.19270730800000002, + "including_init_synchronized_e2e_ms_per_call": 0.21321064400000003, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.417928, + "synchronized_e2e_ms": 7.444712, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 51.884728, + "median": 0.06752, + "min": 0.058784, + "p90": 0.08981760000000008, + "sample_count": 3039 + }, + "host_enqueue_ms": { + "max": 62.434944, + "median": 0.149888, + "min": 0.12224, + "p90": 0.21928400000000023, + "sample_count": 3039 + }, + "sample_count": 3039, + "synchronized_e2e_ms": { + "max": 67.59655, + "median": 0.170305, + "min": 0.143937, + "p90": 0.24613780000000007, + "sample_count": 3039 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 4838, + "candidate_precomputed_gpu_span_ms": 0.018112, + "candidate_precomputed_host_enqueue_ms": 0.052704, + "candidate_precomputed_inter_kernel_gap_ms": 0.00288, + "candidate_precomputed_kernel_sum_ms": 0.015232, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.052704, + "candidate_precomputed_synchronized_e2e_ms": 0.064416, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.015232 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.048256, + "synchronized_e2e_ms": 0.065184 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.018112 + }, + "host_enqueue_ms": { + "median": 0.052704 + }, + "inter_kernel_gap_ms": { + "median": 0.00288 + }, + "kernel_sum_ms": { + "median": 0.015232 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4838, + "submission_ms": { + "median": 0.052704 + }, + "synchronized_e2e_ms": { + "median": 0.064416 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc017028a0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc017007a0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.8215547703180213, + "submission": 0.856709168184578, + "synchronized_e2e": 1.2290269498261301 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 3039, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.032992, + "candidate_public_raw_host_enqueue_ms": 0.045152, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.033216, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.045152, + "candidate_public_raw_synchronized_e2e_ms": 0.079169, + "candidate_public_raw_tflops_from_gpu_span": 128.14798836081474, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.032992 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.053216, + "synchronized_e2e_ms": 0.075712 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.032992 + }, + "host_enqueue_ms": { + "median": 0.045152 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.033216 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3039, + "submission_ms": { + "median": 0.045152 + }, + "synchronized_e2e_ms": { + "median": 0.079169 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.461026, + "after_init_synchronized_e2e_ms_per_call": 2.487906, + "including_init_host_enqueue_ms_per_call": 38.299046, + "including_init_synchronized_e2e_ms_per_call": 486.588119, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2867394, + "after_init_synchronized_e2e_ms_per_call": 0.3200427, + "including_init_host_enqueue_ms_per_call": 3.8705413999999996, + "including_init_synchronized_e2e_ms_per_call": 48.730064, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06931073999999998, + "after_init_synchronized_e2e_ms_per_call": 0.10325637, + "including_init_host_enqueue_ms_per_call": 0.42769093999999996, + "including_init_synchronized_e2e_ms_per_call": 4.9442585, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.047567873999999996, + "after_init_synchronized_e2e_ms_per_call": 0.081577737, + "including_init_host_enqueue_ms_per_call": 0.083405894, + "including_init_synchronized_e2e_ms_per_call": 0.56567795, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.461026, + "synchronized_e2e_ms": 2.487906, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.033696, + "median": 0.032992, + "min": 0.032544, + "p90": 0.033152, + "sample_count": 3039 + }, + "host_enqueue_ms": { + "max": 0.461441, + "median": 0.045152, + "min": 0.035648, + "p90": 0.07016960000000012, + "sample_count": 3039 + }, + "sample_count": 3039, + "synchronized_e2e_ms": { + "max": 5.372582, + "median": 0.079169, + "min": 0.071776, + "p90": 0.10106880000000003, + "sample_count": 3039 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.07152, + "submission_ms": 0.07152, + "synchronized_e2e_ms": 0.0984 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.047264, + "submission_ms": 0.047264, + "synchronized_e2e_ms": 0.075776 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.21456, + "submission_ms": 0.21456, + "synchronized_e2e_ms": 0.232416 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.417928, + "submission_ms": 7.417928, + "synchronized_e2e_ms": 7.444712 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.102176, + "submission_ms": 0.102176, + "synchronized_e2e_ms": 0.121312 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.048256, + "submission_ms": 0.048256, + "synchronized_e2e_ms": 0.065184 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.315362, + "submission_ms": 1.315362, + "synchronized_e2e_ms": 1.341378 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.173729, + "submission_ms": 1.173729, + "synchronized_e2e_ms": 1.196961 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.053216, + "submission_ms": 0.053216, + "synchronized_e2e_ms": 0.075712 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.461026, + "submission_ms": 2.461026, + "synchronized_e2e_ms": 2.487906 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.45224, + "evolution_kernel_ms": 0.176256, + "evolution_speedup": 2.5658, + "evolution_tflops": 23.987, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_a2f8_d224_highn_overlap_b2_n6144_k768_d224", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 15754, + "measurement_schedule_sha256": "5530b7661058e7f2db34c659bf410c9c60db88264ca453041ac17e8ee7aef362", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 4838, + "public_pair_count": 3039, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 4838, + "baseline_public_raw": 3039, + "candidate_precomputed": 4838, + "candidate_public_raw": 3039 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 3152 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:adjacent_a2f8_d224_highn_overlap_b2_n6144_k768_d224", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.3957597173144878, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.151157650090313, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.992360643850692, + "including_init_synchronized_e2e_speedup": 0.08852651209101964, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.8050810095027945, + "including_init_synchronized_e2e_speedup": 0.09154244903105402, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.3538409301043606, + "including_init_synchronized_e2e_speedup": 0.12122372646980331, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.1768121246118906, + "including_init_synchronized_e2e_speedup": 0.37691171098325477, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.046556741028128, + "hot_synchronized_e2e_speedup": 2.151157650090313, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 1028224, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_a2f8_d224_highn_overlap_b2_n6144_k768_d224", + "source": "guard_overlap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 288, + "K": 4096, + "N": 640, + "baseline_07cf_adapter_bench_iters": 3743, + "baseline_07cf_adapter_gpu_span_ms": 0.146496, + "baseline_07cf_adapter_host_enqueue_ms": 0.151648, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.036544, + "baseline_07cf_adapter_kernel_sum_ms": 0.109952, + "baseline_07cf_adapter_submission_ms": 0.151648, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.248896, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.109952 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.264417, + "synchronized_e2e_ms": 0.350977 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.146496 + }, + "host_enqueue_ms": { + "median": 0.151648 + }, + "inter_kernel_gap_ms": { + "median": 0.036544 + }, + "kernel_sum_ms": { + "median": 0.109952 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3743, + "submission_ms": { + "median": 0.151648 + }, + "synchronized_e2e_ms": { + "median": 0.248896 + } + }, + "baseline_07cf_precomputed_bench_iters": 4819, + "baseline_07cf_precomputed_gpu_span_ms": 0.121312, + "baseline_07cf_precomputed_host_enqueue_ms": 0.044128, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.121312, + "baseline_07cf_precomputed_submission_ms": 0.044128, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.170848, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.121312 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.054881, + "synchronized_e2e_ms": 0.150977 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.121312 + }, + "host_enqueue_ms": { + "median": 0.044128 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.121312 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4819, + "submission_ms": { + "median": 0.044128 + }, + "synchronized_e2e_ms": { + "median": 0.170848 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.2075969401213398, + "submission": 3.436548223350254, + "synchronized_e2e": 1.4568271211837422 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.113544, + "after_init_synchronized_e2e_ms_per_call": 8.190472, + "including_init_host_enqueue_ms_per_call": 45.047982, + "including_init_synchronized_e2e_ms_per_call": 463.30179100000004, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9478376, + "after_init_synchronized_e2e_ms_per_call": 1.0430536, + "including_init_host_enqueue_ms_per_call": 4.6412814, + "including_init_synchronized_e2e_ms_per_call": 46.5541855, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23126696, + "after_init_synchronized_e2e_ms_per_call": 0.32831176, + "including_init_host_enqueue_ms_per_call": 0.6006113399999999, + "including_init_synchronized_e2e_ms_per_call": 4.87942495, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.159609896, + "after_init_synchronized_e2e_ms_per_call": 0.256837576, + "including_init_host_enqueue_ms_per_call": 0.196544334, + "including_init_synchronized_e2e_ms_per_call": 0.711948895, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.113544, + "synchronized_e2e_ms": 8.190472, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 45.138733, + "median": 0.146496, + "min": 0.13744, + "p90": 0.16821200000000003, + "sample_count": 3743 + }, + "host_enqueue_ms": { + "max": 45.393871, + "median": 0.151648, + "min": 0.13008, + "p90": 0.2040904, + "sample_count": 3743 + }, + "sample_count": 3743, + "synchronized_e2e_ms": { + "max": 56.642843, + "median": 0.248896, + "min": 0.229024, + "p90": 0.2970690000000001, + "sample_count": 3743 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 4819, + "candidate_precomputed_gpu_span_ms": 0.021376, + "candidate_precomputed_host_enqueue_ms": 0.052928, + "candidate_precomputed_inter_kernel_gap_ms": 0.002176, + "candidate_precomputed_kernel_sum_ms": 0.019136, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.052928, + "candidate_precomputed_synchronized_e2e_ms": 0.065344, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.019136 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.058464, + "synchronized_e2e_ms": 0.07792 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.021376 + }, + "host_enqueue_ms": { + "median": 0.052928 + }, + "inter_kernel_gap_ms": { + "median": 0.002176 + }, + "kernel_sum_ms": { + "median": 0.019136 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4819, + "submission_ms": { + "median": 0.052928 + }, + "synchronized_e2e_ms": { + "median": 0.065344 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc0042a4e0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc0042bda0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.2664670658682635, + "submission": 0.939540507859734, + "synchronized_e2e": 1.1856023506366307 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 3743, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.027072, + "candidate_public_raw_host_enqueue_ms": 0.049728, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.027072, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.049728, + "candidate_public_raw_synchronized_e2e_ms": 0.077472, + "candidate_public_raw_tflops_from_gpu_span": 55.77531914893617, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.026976 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.279616, + "synchronized_e2e_ms": 0.307392 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.027072 + }, + "host_enqueue_ms": { + "median": 0.049728 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.027072 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3743, + "submission_ms": { + "median": 0.049728 + }, + "synchronized_e2e_ms": { + "median": 0.077472 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.800963, + "after_init_synchronized_e2e_ms_per_call": 2.833123, + "including_init_host_enqueue_ms_per_call": 40.158986, + "including_init_synchronized_e2e_ms_per_call": 40.268746, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.32485149999999996, + "after_init_synchronized_e2e_ms_per_call": 0.3530371, + "including_init_host_enqueue_ms_per_call": 4.0606538, + "including_init_synchronized_e2e_ms_per_call": 4.096599400000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07724035000000001, + "after_init_synchronized_e2e_ms_per_call": 0.10502850999999999, + "including_init_host_enqueue_ms_per_call": 0.45082057999999997, + "including_init_synchronized_e2e_ms_per_call": 0.47938474, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.052479235000000006, + "after_init_synchronized_e2e_ms_per_call": 0.080227651, + "including_init_host_enqueue_ms_per_call": 0.08983725799999999, + "including_init_synchronized_e2e_ms_per_call": 0.117663274, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.800963, + "synchronized_e2e_ms": 2.833123, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.027936, + "median": 0.027072, + "min": 0.026176, + "p90": 0.027456, + "sample_count": 3743 + }, + "host_enqueue_ms": { + "max": 40.699946, + "median": 0.049728, + "min": 0.039072, + "p90": 0.06981760000000001, + "sample_count": 3743 + }, + "sample_count": 3743, + "synchronized_e2e_ms": { + "max": 40.808522, + "median": 0.077472, + "min": 0.06768, + "p90": 0.094784, + "sample_count": 3743 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.076832, + "submission_ms": 0.076832, + "synchronized_e2e_ms": 0.171872 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.054881, + "submission_ms": 0.054881, + "synchronized_e2e_ms": 0.150977 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.264417, + "submission_ms": 0.264417, + "synchronized_e2e_ms": 0.350977 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.113544, + "submission_ms": 8.113544, + "synchronized_e2e_ms": 8.190472 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.076864, + "submission_ms": 0.076864, + "synchronized_e2e_ms": 0.098816 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.058464, + "submission_ms": 0.058464, + "synchronized_e2e_ms": 0.07792 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.472033, + "submission_ms": 1.472033, + "synchronized_e2e_ms": 1.496705 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.561026, + "submission_ms": 1.561026, + "synchronized_e2e_ms": 1.585026 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.279616, + "submission_ms": 0.279616, + "synchronized_e2e_ms": 0.307392 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.800963, + "submission_ms": 2.800963, + "synchronized_e2e_ms": 2.833123 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 57.6586, + "evolution_kernel_ms": 0.24416, + "evolution_speedup": 236.1509, + "evolution_tflops": 6.1843, + "expected_route": "d288_parent_splitk_hybrid_20260629_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_a2f8_d288_splitk_probe_b1_n640_k4096_d288", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 17124, + "measurement_schedule_sha256": "6e615ca56d750fdf1c342628993e575b2343600e10247cce35256ecefd10a240", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 4819, + "public_pair_count": 3743, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 4819, + "baseline_public_raw": 3743, + "candidate_precomputed": 4819, + "candidate_public_raw": 3743 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 3426 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:adjacent_a2f8_d288_splitk_probe_b1_n640_k4096_d288", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 5.675149700598803, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 3.2127220156959937, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.890969435495741, + "including_init_synchronized_e2e_speedup": 11.50524506027578, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.95451554525006, + "including_init_synchronized_e2e_speedup": 11.364104945189416, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 3.125929902271298, + "including_init_synchronized_e2e_speedup": 10.178515382029056, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 3.2013597905290787, + "including_init_synchronized_e2e_speedup": 6.050731641208624, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 5.411347517730496, + "hot_synchronized_e2e_speedup": 3.2127220156959937, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 1028288, + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_a2f8_d288_splitk_probe_b1_n640_k4096_d288", + "source": "heldout_neighborhood", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 352, + "K": 256, + "N": 1024, + "baseline_07cf_adapter_bench_iters": 4198, + "baseline_07cf_adapter_gpu_span_ms": 0.054368, + "baseline_07cf_adapter_host_enqueue_ms": 0.14848, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.034112500000000004, + "baseline_07cf_adapter_kernel_sum_ms": 0.020256, + "baseline_07cf_adapter_submission_ms": 0.14848, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.168256, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.020256 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.179712, + "synchronized_e2e_ms": 0.197088 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.054368 + }, + "host_enqueue_ms": { + "median": 0.14848 + }, + "inter_kernel_gap_ms": { + "median": 0.034112500000000004 + }, + "kernel_sum_ms": { + "median": 0.020256 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 4198, + "submission_ms": { + "median": 0.14848 + }, + "synchronized_e2e_ms": { + "median": 0.168256 + } + }, + "baseline_07cf_precomputed_bench_iters": 7862, + "baseline_07cf_precomputed_gpu_span_ms": 0.01264, + "baseline_07cf_precomputed_host_enqueue_ms": 0.041056, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.01264, + "baseline_07cf_precomputed_submission_ms": 0.041056, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.060448, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.01264 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.039296, + "synchronized_e2e_ms": 0.055328 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.01264 + }, + "host_enqueue_ms": { + "median": 0.041056 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.01264 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 7862, + "submission_ms": { + "median": 0.041056 + }, + "synchronized_e2e_ms": { + "median": 0.060448 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 4.30126582278481, + "submission": 3.6165237724084176, + "synchronized_e2e": 2.7834833245103225 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.90916, + "after_init_synchronized_e2e_ms_per_call": 7.935272, + "including_init_host_enqueue_ms_per_call": 42.137067, + "including_init_synchronized_e2e_ms_per_call": 42.280715, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.924548, + "after_init_synchronized_e2e_ms_per_call": 0.9449576000000001, + "including_init_host_enqueue_ms_per_call": 4.3473387, + "including_init_synchronized_e2e_ms_per_call": 4.3795019, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2260868, + "after_init_synchronized_e2e_ms_per_call": 0.24592616, + "including_init_host_enqueue_ms_per_call": 0.56836587, + "including_init_synchronized_e2e_ms_per_call": 0.58938059, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15624068, + "after_init_synchronized_e2e_ms_per_call": 0.17602301599999998, + "including_init_host_enqueue_ms_per_call": 0.19046858700000002, + "including_init_synchronized_e2e_ms_per_call": 0.21036845899999998, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.90916, + "synchronized_e2e_ms": 7.935272, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.513382, + "median": 0.054368, + "min": 0.044928, + "p90": 0.07683840000000002, + "sample_count": 4198 + }, + "host_enqueue_ms": { + "max": 58.489309, + "median": 0.14848, + "min": 0.121536, + "p90": 0.21753280000000003, + "sample_count": 4198 + }, + "sample_count": 4198, + "synchronized_e2e_ms": { + "max": 58.656989, + "median": 0.168256, + "min": 0.137696, + "p90": 0.24296960000000004, + "sample_count": 4198 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 7862, + "candidate_precomputed_gpu_span_ms": 0.017696, + "candidate_precomputed_host_enqueue_ms": 0.052608, + "candidate_precomputed_inter_kernel_gap_ms": 0.003904, + "candidate_precomputed_kernel_sum_ms": 0.013856, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.052608, + "candidate_precomputed_synchronized_e2e_ms": 0.064672, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.013856 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.034912, + "synchronized_e2e_ms": 0.047712 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.017696 + }, + "host_enqueue_ms": { + "median": 0.052608 + }, + "inter_kernel_gap_ms": { + "median": 0.003904 + }, + "kernel_sum_ms": { + "median": 0.013856 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 7862, + "submission_ms": { + "median": 0.052608 + }, + "synchronized_e2e_ms": { + "median": 0.064672 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04737050", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc0270ea20" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.3652802893309224, + "submission": 0.8534063260340632, + "synchronized_e2e": 1.1004530554181098 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 4198, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.02416, + "candidate_public_raw_host_enqueue_ms": 0.044896, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.024352, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.044896, + "candidate_public_raw_synchronized_e2e_ms": 0.0711685, + "candidate_public_raw_tflops_from_gpu_span": 30.554532450331124, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.02416 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.05024, + "synchronized_e2e_ms": 0.067936 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.02416 + }, + "host_enqueue_ms": { + "median": 0.044896 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.024352 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 4198, + "submission_ms": { + "median": 0.044896 + }, + "synchronized_e2e_ms": { + "median": 0.0711685 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.352803, + "after_init_synchronized_e2e_ms_per_call": 2.375299, + "including_init_host_enqueue_ms_per_call": 36.930343, + "including_init_synchronized_e2e_ms_per_call": 450.46309099999996, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2756867, + "after_init_synchronized_e2e_ms_per_call": 0.30158155, + "including_init_host_enqueue_ms_per_call": 3.7334407, + "including_init_synchronized_e2e_ms_per_call": 45.11036075, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06797507, + "after_init_synchronized_e2e_ms_per_call": 0.094209805, + "including_init_host_enqueue_ms_per_call": 0.41375047000000004, + "including_init_synchronized_e2e_ms_per_call": 4.5750877249999995, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.047203907, + "after_init_synchronized_e2e_ms_per_call": 0.0734726305, + "including_init_host_enqueue_ms_per_call": 0.081781447, + "including_init_synchronized_e2e_ms_per_call": 0.5215604225, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.352803, + "synchronized_e2e_ms": 2.375299, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.024736, + "median": 0.02416, + "min": 0.023456, + "p90": 0.024384, + "sample_count": 4198 + }, + "host_enqueue_ms": { + "max": 62.678849, + "median": 0.044896, + "min": 0.035776, + "p90": 0.06812480000000001, + "sample_count": 4198 + }, + "sample_count": 4198, + "synchronized_e2e_ms": { + "max": 62.857858, + "median": 0.0711685, + "min": 0.06272, + "p90": 0.091424, + "sample_count": 4198 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.062273, + "submission_ms": 0.062273, + "synchronized_e2e_ms": 0.077697 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.039296, + "submission_ms": 0.039296, + "synchronized_e2e_ms": 0.055328 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.179712, + "submission_ms": 0.179712, + "synchronized_e2e_ms": 0.197088 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.90916, + "submission_ms": 7.90916, + "synchronized_e2e_ms": 7.935272 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.065152, + "submission_ms": 0.065152, + "synchronized_e2e_ms": 0.079776 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.034912, + "submission_ms": 0.034912, + "synchronized_e2e_ms": 0.047712 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.184449, + "submission_ms": 1.184449, + "synchronized_e2e_ms": 1.210433 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.048353, + "submission_ms": 1.048353, + "synchronized_e2e_ms": 1.069601 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.05024, + "submission_ms": 0.05024, + "synchronized_e2e_ms": 0.067936 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.352803, + "submission_ms": 2.352803, + "synchronized_e2e_ms": 2.375299 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.479503, + "evolution_kernel_ms": 0.176096, + "evolution_speedup": 2.723, + "evolution_tflops": 4.192, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_a2f8_d352_smallk_boundary_b4_n1024_k256_d352", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 24120, + "measurement_schedule_sha256": "2174df98b8174412f703abc891f88f4e8847da0b0882d4a46be722e2bcdae1c3", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 7862, + "public_pair_count": 4198, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 7862, + "baseline_public_raw": 4198, + "candidate_precomputed": 7862, + "candidate_public_raw": 4198 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 4826 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:adjacent_a2f8_d352_smallk_boundary_b4_n1024_k256_d352", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.7142857142857143, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.364192023156312, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.3407465754837604, + "including_init_synchronized_e2e_speedup": 0.09386055338327376, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 3.133340219254129, + "including_init_synchronized_e2e_speedup": 0.09708416929474456, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.6104093942238817, + "including_init_synchronized_e2e_speedup": 0.1288238882894863, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.395763086228415, + "including_init_synchronized_e2e_speedup": 0.40334436802478046, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.2503311258278145, + "hot_synchronized_e2e_speedup": 2.364192023156312, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 1028352, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_a2f8_d352_smallk_boundary_b4_n1024_k256_d352", + "source": "guard_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 416, + "K": 768, + "N": 2560, + "baseline_07cf_adapter_bench_iters": 2488, + "baseline_07cf_adapter_gpu_span_ms": 0.06976, + "baseline_07cf_adapter_host_enqueue_ms": 0.150768, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.03312, + "baseline_07cf_adapter_kernel_sum_ms": 0.036672, + "baseline_07cf_adapter_submission_ms": 0.150768, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.17292849999999999, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.036672 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.22784, + "synchronized_e2e_ms": 0.250464 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.06976 + }, + "host_enqueue_ms": { + "median": 0.150768 + }, + "inter_kernel_gap_ms": { + "median": 0.03312 + }, + "kernel_sum_ms": { + "median": 0.036672 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2488, + "submission_ms": { + "median": 0.150768 + }, + "synchronized_e2e_ms": { + "median": 0.17292849999999999 + } + }, + "baseline_07cf_precomputed_bench_iters": 3498, + "baseline_07cf_precomputed_gpu_span_ms": 0.028544, + "baseline_07cf_precomputed_host_enqueue_ms": 0.044256, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.028544, + "baseline_07cf_precomputed_submission_ms": 0.044256, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.078816, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.028544 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.046144, + "synchronized_e2e_ms": 0.07296 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.028544 + }, + "host_enqueue_ms": { + "median": 0.044256 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.028544 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3498, + "submission_ms": { + "median": 0.044256 + }, + "synchronized_e2e_ms": { + "median": 0.078816 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.4439461883408073, + "submission": 3.4067245119305865, + "synchronized_e2e": 2.1940786134794963 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 6.929543, + "after_init_synchronized_e2e_ms_per_call": 6.955463, + "including_init_host_enqueue_ms_per_call": 41.378987, + "including_init_synchronized_e2e_ms_per_call": 490.996957, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8286455, + "after_init_synchronized_e2e_ms_per_call": 0.85118195, + "including_init_host_enqueue_ms_per_call": 4.2735899, + "including_init_synchronized_e2e_ms_per_call": 49.25533135, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.21855575000000002, + "after_init_synchronized_e2e_ms_per_call": 0.240753845, + "including_init_host_enqueue_ms_per_call": 0.56305019, + "including_init_synchronized_e2e_ms_per_call": 5.081168785, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.157546775, + "after_init_synchronized_e2e_ms_per_call": 0.17971103449999998, + "including_init_host_enqueue_ms_per_call": 0.191996219, + "including_init_synchronized_e2e_ms_per_call": 0.6637525284999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 6.929543, + "synchronized_e2e_ms": 6.955463, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.75616, + "median": 0.06976, + "min": 0.061152, + "p90": 0.0989123, + "sample_count": 2488 + }, + "host_enqueue_ms": { + "max": 1.328834, + "median": 0.150768, + "min": 0.126912, + "p90": 0.3017440000000008, + "sample_count": 2488 + }, + "sample_count": 2488, + "synchronized_e2e_ms": { + "max": 1.385026, + "median": 0.17292849999999999, + "min": 0.149184, + "p90": 0.3439587000000005, + "sample_count": 2488 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3498, + "candidate_precomputed_gpu_span_ms": 0.031488, + "candidate_precomputed_host_enqueue_ms": 0.055936, + "candidate_precomputed_inter_kernel_gap_ms": 0.003872, + "candidate_precomputed_kernel_sum_ms": 0.027584, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.055936, + "candidate_precomputed_synchronized_e2e_ms": 0.068832, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.027584 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.044576, + "synchronized_e2e_ms": 0.061184 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.031488 + }, + "host_enqueue_ms": { + "median": 0.055936 + }, + "inter_kernel_gap_ms": { + "median": 0.003872 + }, + "kernel_sum_ms": { + "median": 0.027584 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3498, + "submission_ms": { + "median": 0.055936 + }, + "synchronized_e2e_ms": { + "median": 0.068832 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7dcd6f30", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7dcd49e0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.2754065040650406, + "submission": 0.8318077803203662, + "synchronized_e2e": 1.284053928405393 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2488, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.04016, + "candidate_public_raw_host_enqueue_ms": 0.046528, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.040384, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.046528, + "candidate_public_raw_synchronized_e2e_ms": 0.088384, + "candidate_public_raw_tflops_from_gpu_span": 81.46307569721115, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.04016 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.070592, + "synchronized_e2e_ms": 0.096672 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.04016 + }, + "host_enqueue_ms": { + "median": 0.046528 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.040384 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2488, + "submission_ms": { + "median": 0.046528 + }, + "synchronized_e2e_ms": { + "median": 0.088384 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.276195, + "after_init_synchronized_e2e_ms_per_call": 2.301571, + "including_init_host_enqueue_ms_per_call": 37.150279, + "including_init_synchronized_e2e_ms_per_call": 37.244647, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2694947, + "after_init_synchronized_e2e_ms_per_call": 0.3097027, + "including_init_host_enqueue_ms_per_call": 3.7569030999999997, + "including_init_synchronized_e2e_ms_per_call": 3.8040103000000003, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06882467, + "after_init_synchronized_e2e_ms_per_call": 0.11051587000000002, + "including_init_host_enqueue_ms_per_call": 0.41756551, + "including_init_synchronized_e2e_ms_per_call": 0.45994663, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.048757667, + "after_init_synchronized_e2e_ms_per_call": 0.09059718700000001, + "including_init_host_enqueue_ms_per_call": 0.08363175099999999, + "including_init_synchronized_e2e_ms_per_call": 0.125540263, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.276195, + "synchronized_e2e_ms": 2.301571, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.040704, + "median": 0.04016, + "min": 0.039712, + "p90": 0.04032, + "sample_count": 2488 + }, + "host_enqueue_ms": { + "max": 63.193058, + "median": 0.046528, + "min": 0.037184, + "p90": 0.11134720000000013, + "sample_count": 2488 + }, + "sample_count": 2488, + "synchronized_e2e_ms": { + "max": 68.373447, + "median": 0.088384, + "min": 0.07984, + "p90": 0.14460550000000005, + "sample_count": 2488 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.065344, + "submission_ms": 0.065344, + "synchronized_e2e_ms": 0.090592 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.046144, + "submission_ms": 0.046144, + "synchronized_e2e_ms": 0.07296 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.22784, + "submission_ms": 0.22784, + "synchronized_e2e_ms": 0.250464 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 6.929543, + "submission_ms": 6.929543, + "synchronized_e2e_ms": 6.955463 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.070528, + "submission_ms": 0.070528, + "synchronized_e2e_ms": 0.088672 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.044576, + "submission_ms": 0.044576, + "synchronized_e2e_ms": 0.061184 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.206594, + "submission_ms": 1.206594, + "synchronized_e2e_ms": 1.230178 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.061537, + "submission_ms": 1.061537, + "synchronized_e2e_ms": 1.084801 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.070592, + "submission_ms": 0.070592, + "synchronized_e2e_ms": 0.096672 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.276195, + "submission_ms": 2.276195, + "synchronized_e2e_ms": 2.301571 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.492399, + "evolution_kernel_ms": 0.189631, + "evolution_speedup": 2.5966, + "evolution_tflops": 17.2522, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_a2f8_d416_random_b2_n2560_k768_d416", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 11972, + "measurement_schedule_sha256": "fcb159c17a6e30dff234f63ba879bae41316c39286a9bf7fd64b43cabe654b9f", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3498, + "public_pair_count": 2488, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3498, + "baseline_public_raw": 2488, + "candidate_precomputed": 3498, + "candidate_public_raw": 2488 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2396 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:adjacent_a2f8_d416_random_b2_n2560_k768_d416", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.9065040650406504, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.9565588794351916, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.0220501561759336, + "including_init_synchronized_e2e_speedup": 13.183020824442234, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.7483840147341305, + "including_init_synchronized_e2e_speedup": 12.948264453963228, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.1784549585502964, + "including_init_synchronized_e2e_speedup": 11.047300824880486, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.983627090982416, + "including_init_synchronized_e2e_speedup": 5.287168535723076, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.7370517928286853, + "hot_synchronized_e2e_speedup": 1.9565588794351916, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 1028416, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_a2f8_d416_random_b2_n2560_k768_d416", + "source": "random_legal", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 480, + "K": 4096, + "N": 896, + "baseline_07cf_adapter_bench_iters": 2310, + "baseline_07cf_adapter_gpu_span_ms": 0.155392, + "baseline_07cf_adapter_host_enqueue_ms": 0.14966449999999998, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.033888, + "baseline_07cf_adapter_kernel_sum_ms": 0.121536, + "baseline_07cf_adapter_submission_ms": 0.14966449999999998, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.257168, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.121536 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.187168, + "synchronized_e2e_ms": 0.289536 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.155392 + }, + "host_enqueue_ms": { + "median": 0.14966449999999998 + }, + "inter_kernel_gap_ms": { + "median": 0.033888 + }, + "kernel_sum_ms": { + "median": 0.121536 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2310, + "submission_ms": { + "median": 0.14966449999999998 + }, + "synchronized_e2e_ms": { + "median": 0.257168 + } + }, + "baseline_07cf_precomputed_bench_iters": 2771, + "baseline_07cf_precomputed_gpu_span_ms": 0.127585, + "baseline_07cf_precomputed_host_enqueue_ms": 0.043808, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.127585, + "baseline_07cf_precomputed_submission_ms": 0.043808, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.176833, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.127585 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.048, + "synchronized_e2e_ms": 0.153056 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.127585 + }, + "host_enqueue_ms": { + "median": 0.043808 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.127585 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2771, + "submission_ms": { + "median": 0.043808 + }, + "synchronized_e2e_ms": { + "median": 0.176833 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.2179488184347689, + "submission": 3.4163737216946672, + "synchronized_e2e": 1.4542986885931926 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.670152, + "after_init_synchronized_e2e_ms_per_call": 7.75924, + "including_init_host_enqueue_ms_per_call": 43.22142, + "including_init_synchronized_e2e_ms_per_call": 43.390477, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9017132499999999, + "after_init_synchronized_e2e_ms_per_call": 1.0073752, + "including_init_host_enqueue_ms_per_call": 4.45684005, + "including_init_synchronized_e2e_ms_per_call": 4.5704989, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22486937499999995, + "after_init_synchronized_e2e_ms_per_call": 0.33218872, + "including_init_host_enqueue_ms_per_call": 0.5803820550000001, + "including_init_synchronized_e2e_ms_per_call": 0.68850109, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15718498749999998, + "after_init_synchronized_e2e_ms_per_call": 0.264670072, + "including_init_host_enqueue_ms_per_call": 0.19273625549999998, + "including_init_synchronized_e2e_ms_per_call": 0.300301309, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.670152, + "synchronized_e2e_ms": 7.75924, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.371718, + "median": 0.155392, + "min": 0.145728, + "p90": 0.17347839999999998, + "sample_count": 2310 + }, + "host_enqueue_ms": { + "max": 96.058435, + "median": 0.14966449999999998, + "min": 0.123616, + "p90": 0.19866559999999997, + "sample_count": 2310 + }, + "sample_count": 2310, + "synchronized_e2e_ms": { + "max": 101.228585, + "median": 0.257168, + "min": 0.233024, + "p90": 0.30084809999999995, + "sample_count": 2310 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 2771, + "candidate_precomputed_gpu_span_ms": 0.035936, + "candidate_precomputed_host_enqueue_ms": 0.050241, + "candidate_precomputed_inter_kernel_gap_ms": 0.00208, + "candidate_precomputed_kernel_sum_ms": 0.033824, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.050241, + "candidate_precomputed_synchronized_e2e_ms": 0.069952, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.033824 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.041312, + "synchronized_e2e_ms": 0.056609 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.035936 + }, + "host_enqueue_ms": { + "median": 0.050241 + }, + "inter_kernel_gap_ms": { + "median": 0.00208 + }, + "kernel_sum_ms": { + "median": 0.033824 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2771, + "submission_ms": { + "median": 0.050241 + }, + "synchronized_e2e_ms": { + "median": 0.069952 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb937a0950", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb937a3320" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.2061581700801425, + "submission": 0.9630381560876575, + "synchronized_e2e": 1.3218278247941446 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2310, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.0433445, + "candidate_public_raw_host_enqueue_ms": 0.048384, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.043584, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.048384, + "candidate_public_raw_synchronized_e2e_ms": 0.0924645, + "candidate_public_raw_tflops_from_gpu_span": 81.28402357853936, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.0433445 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.240385, + "synchronized_e2e_ms": 0.266721 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.0433445 + }, + "host_enqueue_ms": { + "median": 0.048384 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.043584 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2310, + "submission_ms": { + "median": 0.048384 + }, + "synchronized_e2e_ms": { + "median": 0.0924645 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.31869, + "after_init_synchronized_e2e_ms_per_call": 2.34333, + "including_init_host_enqueue_ms_per_call": 38.156710000000004, + "including_init_synchronized_e2e_ms_per_call": 486.443543, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2754146, + "after_init_synchronized_e2e_ms_per_call": 0.31755104999999995, + "including_init_host_enqueue_ms_per_call": 3.8592166000000008, + "including_init_synchronized_e2e_ms_per_call": 48.727572349999996, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07108706000000001, + "after_init_synchronized_e2e_ms_per_call": 0.11497315500000001, + "including_init_host_enqueue_ms_per_call": 0.4294672600000001, + "including_init_synchronized_e2e_ms_per_call": 4.955975284999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.050654306, + "after_init_synchronized_e2e_ms_per_call": 0.09471536550000001, + "including_init_host_enqueue_ms_per_call": 0.08649232600000001, + "including_init_synchronized_e2e_ms_per_call": 0.5788155785, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.31869, + "synchronized_e2e_ms": 2.34333, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.044416, + "median": 0.0433445, + "min": 0.042016, + "p90": 0.043776, + "sample_count": 2310 + }, + "host_enqueue_ms": { + "max": 15.679792, + "median": 0.048384, + "min": 0.037952, + "p90": 0.0669152, + "sample_count": 2310 + }, + "sample_count": 2310, + "synchronized_e2e_ms": { + "max": 15.713616, + "median": 0.0924645, + "min": 0.08368, + "p90": 0.1087104, + "sample_count": 2310 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.084768, + "submission_ms": 0.084768, + "synchronized_e2e_ms": 0.187488 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.048, + "submission_ms": 0.048, + "synchronized_e2e_ms": 0.153056 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.187168, + "submission_ms": 0.187168, + "synchronized_e2e_ms": 0.289536 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.670152, + "submission_ms": 7.670152, + "synchronized_e2e_ms": 7.75924 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.064128, + "submission_ms": 0.064128, + "synchronized_e2e_ms": 0.0816 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.041312, + "submission_ms": 0.041312, + "synchronized_e2e_ms": 0.056609 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.384033, + "submission_ms": 1.384033, + "synchronized_e2e_ms": 1.404801 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.170977, + "submission_ms": 1.170977, + "synchronized_e2e_ms": 1.192289 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.240385, + "submission_ms": 0.240385, + "synchronized_e2e_ms": 0.266721 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.31869, + "submission_ms": 2.31869, + "synchronized_e2e_ms": 2.34333 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 32.096809, + "evolution_kernel_ms": 0.274432, + "evolution_speedup": 116.9572, + "evolution_tflops": 12.8382, + "expected_route": "d480_splitk_k1024_eac2_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_a2f8_d480_highk_request_b1_n896_k4096_d480", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 10162, + "measurement_schedule_sha256": "e77ec20b3385598fb086b734d4438a2c3aebee9626f3e40724171fbcd65c07c5", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 2771, + "public_pair_count": 2310, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 2771, + "baseline_public_raw": 2310, + "candidate_precomputed": 2771, + "candidate_public_raw": 2310 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2034 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:adjacent_a2f8_d480_highk_request_b1_n896_k4096_d480", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 3.550339492430988, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.781261997847823, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.3112024341428654, + "including_init_synchronized_e2e_speedup": 0.08919940993029071, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 3.1723252056637827, + "including_init_synchronized_e2e_speedup": 0.09379697529708762, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.889272021803698, + "including_init_synchronized_e2e_speedup": 0.13892343089035403, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.7943731262906857, + "including_init_synchronized_e2e_speedup": 0.5188203637818984, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 3.5850453921489462, + "hot_synchronized_e2e_speedup": 2.781261997847823, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 1028480, + "selected_route": "d480_splitk_k1024_eac2_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_a2f8_d480_highk_request_b1_n896_k4096_d480", + "source": "request_specific", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 48, + "K": 1024, + "N": 768, + "baseline_07cf_adapter_bench_iters": 6714, + "baseline_07cf_adapter_gpu_span_ms": 0.0615205, + "baseline_07cf_adapter_host_enqueue_ms": 0.1549925, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.041408, + "baseline_07cf_adapter_kernel_sum_ms": 0.02016, + "baseline_07cf_adapter_submission_ms": 0.1549925, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.175008, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.02016 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.183136, + "synchronized_e2e_ms": 0.201632 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.0615205 + }, + "host_enqueue_ms": { + "median": 0.1549925 + }, + "inter_kernel_gap_ms": { + "median": 0.041408 + }, + "kernel_sum_ms": { + "median": 0.02016 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 6714, + "submission_ms": { + "median": 0.1549925 + }, + "synchronized_e2e_ms": { + "median": 0.175008 + } + }, + "baseline_07cf_precomputed_bench_iters": 5398, + "baseline_07cf_precomputed_gpu_span_ms": 0.02, + "baseline_07cf_precomputed_host_enqueue_ms": 0.04496, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.02, + "baseline_07cf_precomputed_submission_ms": 0.04496, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.070144, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.02 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.040288, + "synchronized_e2e_ms": 0.06 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.02 + }, + "host_enqueue_ms": { + "median": 0.04496 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.02 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5398, + "submission_ms": { + "median": 0.04496 + }, + "synchronized_e2e_ms": { + "median": 0.070144 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 3.076025, + "submission": 3.447342081850534, + "synchronized_e2e": 2.4949817518248176 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.766248, + "after_init_synchronized_e2e_ms_per_call": 7.788968, + "including_init_host_enqueue_ms_per_call": 44.700686, + "including_init_synchronized_e2e_ms_per_call": 462.90028700000005, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.91611805, + "after_init_synchronized_e2e_ms_per_call": 0.9364039999999999, + "including_init_host_enqueue_ms_per_call": 4.60956185, + "including_init_synchronized_e2e_ms_per_call": 46.447535900000005, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23110505500000003, + "after_init_synchronized_e2e_ms_per_call": 0.2511476, + "including_init_host_enqueue_ms_per_call": 0.600449435, + "including_init_synchronized_e2e_ms_per_call": 4.80226079, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.1626037555, + "after_init_synchronized_e2e_ms_per_call": 0.18262196, + "including_init_host_enqueue_ms_per_call": 0.19953819350000002, + "including_init_synchronized_e2e_ms_per_call": 0.637733279, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.766248, + "synchronized_e2e_ms": 7.788968, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 10.411947, + "median": 0.0615205, + "min": 0.050976, + "p90": 0.08990079999999999, + "sample_count": 6714 + }, + "host_enqueue_ms": { + "max": 72.685227, + "median": 0.1549925, + "min": 0.118369, + "p90": 0.32832069999999997, + "sample_count": 6714 + }, + "sample_count": 6714, + "synchronized_e2e_ms": { + "max": 72.858795, + "median": 0.175008, + "min": 0.135297, + "p90": 0.3703302999999995, + "sample_count": 6714 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 5398, + "candidate_precomputed_gpu_span_ms": 0.018016, + "candidate_precomputed_host_enqueue_ms": 0.059264, + "candidate_precomputed_inter_kernel_gap_ms": 0.006272, + "candidate_precomputed_kernel_sum_ms": 0.011712, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.059264, + "candidate_precomputed_synchronized_e2e_ms": 0.072096, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.011712 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.04128, + "synchronized_e2e_ms": 0.055488 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.018016 + }, + "host_enqueue_ms": { + "median": 0.059264 + }, + "inter_kernel_gap_ms": { + "median": 0.006272 + }, + "kernel_sum_ms": { + "median": 0.011712 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5398, + "submission_ms": { + "median": 0.059264 + }, + "synchronized_e2e_ms": { + "median": 0.072096 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb67415460", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb67414ad0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 0.8401420959147424, + "submission": 0.8012958963282938, + "synchronized_e2e": 0.8845983133599645 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 6714, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.015136, + "candidate_public_raw_host_enqueue_ms": 0.047488, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.015296, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.047488, + "candidate_public_raw_synchronized_e2e_ms": 0.063776, + "candidate_public_raw_tflops_from_gpu_span": 9.975881606765327, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.015136 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.06992, + "synchronized_e2e_ms": 0.0888 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.015136 + }, + "host_enqueue_ms": { + "median": 0.047488 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.015296 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 6714, + "submission_ms": { + "median": 0.047488 + }, + "synchronized_e2e_ms": { + "median": 0.063776 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.431203, + "after_init_synchronized_e2e_ms_per_call": 2.455459, + "including_init_host_enqueue_ms_per_call": 39.789226, + "including_init_synchronized_e2e_ms_per_call": 39.891082, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.28585950000000004, + "after_init_synchronized_e2e_ms_per_call": 0.30294429999999994, + "including_init_host_enqueue_ms_per_call": 4.0216617999999995, + "including_init_synchronized_e2e_ms_per_call": 4.0465066, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07132515, + "after_init_synchronized_e2e_ms_per_call": 0.08769283, + "including_init_host_enqueue_ms_per_call": 0.44490538, + "including_init_synchronized_e2e_ms_per_call": 0.46204905999999996, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.04987171500000001, + "after_init_synchronized_e2e_ms_per_call": 0.06616768299999999, + "including_init_host_enqueue_ms_per_call": 0.087229738, + "including_init_synchronized_e2e_ms_per_call": 0.103603306, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.431203, + "synchronized_e2e_ms": 2.455459, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.015648, + "median": 0.015136, + "min": 0.01472, + "p90": 0.015296, + "sample_count": 6714 + }, + "host_enqueue_ms": { + "max": 32.939138, + "median": 0.047488, + "min": 0.033345, + "p90": 0.12232959999999986, + "sample_count": 6714 + }, + "sample_count": 6714, + "synchronized_e2e_ms": { + "max": 70.774761, + "median": 0.063776, + "min": 0.051264, + "p90": 0.15982399999999986, + "sample_count": 6714 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.063872, + "submission_ms": 0.063872, + "synchronized_e2e_ms": 0.082304 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.040288, + "submission_ms": 0.040288, + "synchronized_e2e_ms": 0.06 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.183136, + "submission_ms": 0.183136, + "synchronized_e2e_ms": 0.201632 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.766248, + "submission_ms": 7.766248, + "synchronized_e2e_ms": 7.788968 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.070048, + "submission_ms": 0.070048, + "synchronized_e2e_ms": 0.085408 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.04128, + "submission_ms": 0.04128, + "synchronized_e2e_ms": 0.055488 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.317761, + "submission_ms": 1.317761, + "synchronized_e2e_ms": 1.340769 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.12957, + "submission_ms": 1.12957, + "synchronized_e2e_ms": 1.152354 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.06992, + "submission_ms": 0.06992, + "synchronized_e2e_ms": 0.0888 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.431203, + "submission_ms": 2.431203, + "synchronized_e2e_ms": 2.455459 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.497055, + "evolution_kernel_ms": 0.173184, + "evolution_speedup": 2.8701, + "evolution_tflops": 0.8719, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_a2f8_d48_k1024_boundary_b2_n768_k1024_d48", + "measurement_order": [ + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 24224, + "measurement_schedule_sha256": "f2642e6937ffa29f4736c93f6a264e23e9446f47fadff068fdf1c4e503deda18", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 5398, + "public_pair_count": 6714, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 5398, + "baseline_public_raw": 6714, + "candidate_precomputed": 5398, + "candidate_public_raw": 6714 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 4846 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:adjacent_a2f8_d48_k1024_boundary_b2_n768_k1024_d48", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.1101243339253997, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.7441043652784747, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.172102649647174, + "including_init_synchronized_e2e_speedup": 11.604104571543084, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 3.0910104596785617, + "including_init_synchronized_e2e_speedup": 11.478428306529887, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.8639468015800156, + "including_init_synchronized_e2e_speedup": 10.393400194343, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.759987228206253, + "including_init_synchronized_e2e_speedup": 6.155530200937796, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 4.064515063424947, + "hot_synchronized_e2e_speedup": 2.7441043652784747, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 1028481, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_a2f8_d48_k1024_boundary_b2_n768_k1024_d48", + "source": "guard_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 64, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 112, + "K": 768, + "N": 3200, + "baseline_07cf_adapter_bench_iters": 5431, + "baseline_07cf_adapter_gpu_span_ms": 0.058208, + "baseline_07cf_adapter_host_enqueue_ms": 0.142016, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.03568, + "baseline_07cf_adapter_kernel_sum_ms": 0.022528, + "baseline_07cf_adapter_submission_ms": 0.142016, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.16128, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.022528 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.198336, + "synchronized_e2e_ms": 0.216 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.058208 + }, + "host_enqueue_ms": { + "median": 0.142016 + }, + "inter_kernel_gap_ms": { + "median": 0.03568 + }, + "kernel_sum_ms": { + "median": 0.022528 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 5431, + "submission_ms": { + "median": 0.142016 + }, + "synchronized_e2e_ms": { + "median": 0.16128 + } + }, + "baseline_07cf_precomputed_bench_iters": 5288, + "baseline_07cf_precomputed_gpu_span_ms": 0.01888, + "baseline_07cf_precomputed_host_enqueue_ms": 0.041696, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.01888, + "baseline_07cf_precomputed_submission_ms": 0.041696, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.067104, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.01888 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.043072, + "synchronized_e2e_ms": 0.063392 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.01888 + }, + "host_enqueue_ms": { + "median": 0.041696 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.01888 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5288, + "submission_ms": { + "median": 0.041696 + }, + "synchronized_e2e_ms": { + "median": 0.067104 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 3.0830508474576273, + "submission": 3.40598618572525, + "synchronized_e2e": 2.40343347639485 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.645224, + "after_init_synchronized_e2e_ms_per_call": 7.672616, + "including_init_host_enqueue_ms_per_call": 41.873131, + "including_init_synchronized_e2e_ms_per_call": 42.018059, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8923368, + "after_init_synchronized_e2e_ms_per_call": 0.9124136, + "including_init_host_enqueue_ms_per_call": 4.3151275, + "including_init_synchronized_e2e_ms_per_call": 4.3469579000000005, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.21704808, + "after_init_synchronized_e2e_ms_per_call": 0.23639336, + "including_init_host_enqueue_ms_per_call": 0.55932715, + "including_init_synchronized_e2e_ms_per_call": 0.5798477900000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.14951920800000001, + "after_init_synchronized_e2e_ms_per_call": 0.16879133600000001, + "including_init_host_enqueue_ms_per_call": 0.18374711500000002, + "including_init_synchronized_e2e_ms_per_call": 0.203136779, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.645224, + "synchronized_e2e_ms": 7.672616, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 50.665749, + "median": 0.058208, + "min": 0.050336, + "p90": 0.079168, + "sample_count": 5431 + }, + "host_enqueue_ms": { + "max": 51.587286, + "median": 0.142016, + "min": 0.118048, + "p90": 0.181312, + "sample_count": 5431 + }, + "sample_count": 5431, + "synchronized_e2e_ms": { + "max": 51.756598, + "median": 0.16128, + "min": 0.134112, + "p90": 0.205473, + "sample_count": 5431 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 5288, + "candidate_precomputed_gpu_span_ms": 0.016288, + "candidate_precomputed_host_enqueue_ms": 0.053792, + "candidate_precomputed_inter_kernel_gap_ms": 0.004704, + "candidate_precomputed_kernel_sum_ms": 0.011616, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.053792, + "candidate_precomputed_synchronized_e2e_ms": 0.065984, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.011616 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.039424, + "synchronized_e2e_ms": 0.053056 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.016288 + }, + "host_enqueue_ms": { + "median": 0.053792 + }, + "inter_kernel_gap_ms": { + "median": 0.004704 + }, + "kernel_sum_ms": { + "median": 0.011616 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5288, + "submission_ms": { + "median": 0.053792 + }, + "synchronized_e2e_ms": { + "median": 0.065984 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc045346b0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04534530" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.1277013752455796, + "submission": 0.810232004759072, + "synchronized_e2e": 0.9733268671193017 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 5431, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.018368, + "candidate_public_raw_host_enqueue_ms": 0.043584, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.01856, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.043584, + "candidate_public_raw_synchronized_e2e_ms": 0.064224, + "candidate_public_raw_tflops_from_gpu_span": 59.94146341463415, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.018368 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.078368, + "synchronized_e2e_ms": 0.096768 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.018368 + }, + "host_enqueue_ms": { + "median": 0.043584 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.01856 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 5431, + "submission_ms": { + "median": 0.043584 + }, + "synchronized_e2e_ms": { + "median": 0.064224 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.530691, + "after_init_synchronized_e2e_ms_per_call": 2.556323, + "including_init_host_enqueue_ms_per_call": 37.108230999999996, + "including_init_synchronized_e2e_ms_per_call": 450.644115, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2922947, + "after_init_synchronized_e2e_ms_per_call": 0.3134339, + "including_init_host_enqueue_ms_per_call": 3.7500487, + "including_init_synchronized_e2e_ms_per_call": 45.122213099999996, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06845506999999999, + "after_init_synchronized_e2e_ms_per_call": 0.08914499, + "including_init_host_enqueue_ms_per_call": 0.41423046999999996, + "including_init_synchronized_e2e_ms_per_call": 4.5700229100000005, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.046071107, + "after_init_synchronized_e2e_ms_per_call": 0.06671609900000001, + "including_init_host_enqueue_ms_per_call": 0.08064864699999999, + "including_init_synchronized_e2e_ms_per_call": 0.514803891, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.530691, + "synchronized_e2e_ms": 2.556323, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.019104, + "median": 0.018368, + "min": 0.017952, + "p90": 0.018592, + "sample_count": 5431 + }, + "host_enqueue_ms": { + "max": 25.683546, + "median": 0.043584, + "min": 0.033568, + "p90": 0.062304, + "sample_count": 5431 + }, + "sample_count": 5431, + "synchronized_e2e_ms": { + "max": 35.652549, + "median": 0.064224, + "min": 0.055488, + "p90": 0.082944, + "sample_count": 5431 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.06656, + "submission_ms": 0.06656, + "synchronized_e2e_ms": 0.085472 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.043072, + "submission_ms": 0.043072, + "synchronized_e2e_ms": 0.063392 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.198336, + "submission_ms": 0.198336, + "synchronized_e2e_ms": 0.216 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.645224, + "submission_ms": 7.645224, + "synchronized_e2e_ms": 7.672616 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.071424, + "submission_ms": 0.071424, + "synchronized_e2e_ms": 0.0864 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.039424, + "submission_ms": 0.039424, + "synchronized_e2e_ms": 0.053056 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.341985, + "submission_ms": 1.341985, + "synchronized_e2e_ms": 1.369025 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.125186, + "submission_ms": 1.125186, + "synchronized_e2e_ms": 1.14893 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.078368, + "submission_ms": 0.078368, + "synchronized_e2e_ms": 0.096768 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.530691, + "submission_ms": 2.530691, + "synchronized_e2e_ms": 2.556323 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.321696, + "evolution_kernel_ms": 0.173088, + "evolution_speedup": 1.8586, + "evolution_tflops": 6.361, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_c44f_d112_odd_tail_b2_n3200_k768_d112", + "measurement_order": [ + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 21438, + "measurement_schedule_sha256": "33a272f93d1b326b383a6f92dc60ce53313f1ac87743dfafcc7456c9f353db02", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 5288, + "public_pair_count": 5431, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 5288, + "baseline_public_raw": 5431, + "candidate_precomputed": 5288, + "candidate_public_raw": 5431 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 4290 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:adjacent_c44f_d112_odd_tail_b2_n3200_k768_d112", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.1591355599214146, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.5112107623318387, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.001426658524764, + "including_init_synchronized_e2e_speedup": 0.09324000381098953, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.9110239830471434, + "including_init_synchronized_e2e_speedup": 0.09633742676508925, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.651785142384334, + "including_init_synchronized_e2e_speedup": 0.12688071841635473, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.5299940873341527, + "including_init_synchronized_e2e_speedup": 0.39459060537675694, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 3.168989547038328, + "hot_synchronized_e2e_speedup": 2.5112107623318387, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 4411201, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_c44f_d112_odd_tail_b2_n3200_k768_d112", + "source": "tail_divisibility", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 6, + "D": 128, + "K": 512, + "N": 6272, + "baseline_07cf_adapter_bench_iters": 2661, + "baseline_07cf_adapter_gpu_span_ms": 0.064672, + "baseline_07cf_adapter_host_enqueue_ms": 0.156512, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.022177, + "baseline_07cf_adapter_kernel_sum_ms": 0.042528, + "baseline_07cf_adapter_submission_ms": 0.156512, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.175488, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.042528 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.220545, + "synchronized_e2e_ms": 0.239457 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.064672 + }, + "host_enqueue_ms": { + "median": 0.156512 + }, + "inter_kernel_gap_ms": { + "median": 0.022177 + }, + "kernel_sum_ms": { + "median": 0.042528 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2661, + "submission_ms": { + "median": 0.156512 + }, + "synchronized_e2e_ms": { + "median": 0.175488 + } + }, + "baseline_07cf_precomputed_bench_iters": 7319, + "baseline_07cf_precomputed_gpu_span_ms": 0.01824, + "baseline_07cf_precomputed_host_enqueue_ms": 0.043008, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.01824, + "baseline_07cf_precomputed_submission_ms": 0.043008, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.06848, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.01824 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.04544, + "synchronized_e2e_ms": 0.068928 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.01824 + }, + "host_enqueue_ms": { + "median": 0.043008 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.01824 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 7319, + "submission_ms": { + "median": 0.043008 + }, + "synchronized_e2e_ms": { + "median": 0.06848 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 3.545614035087719, + "submission": 3.639136904761905, + "synchronized_e2e": 2.5626168224299066 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.384745, + "after_init_synchronized_e2e_ms_per_call": 8.414121, + "including_init_host_enqueue_ms_per_call": 42.834189, + "including_init_synchronized_e2e_ms_per_call": 492.455615, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9793353, + "after_init_synchronized_e2e_ms_per_call": 0.9993513, + "including_init_host_enqueue_ms_per_call": 4.4242797000000005, + "including_init_synchronized_e2e_ms_per_call": 49.4035007, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23879433000000003, + "after_init_synchronized_e2e_ms_per_call": 0.25787433, + "including_init_host_enqueue_ms_per_call": 0.58328877, + "including_init_synchronized_e2e_ms_per_call": 5.09828927, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.164740233, + "after_init_synchronized_e2e_ms_per_call": 0.183726633, + "including_init_host_enqueue_ms_per_call": 0.199189677, + "including_init_synchronized_e2e_ms_per_call": 0.6677681270000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.384745, + "synchronized_e2e_ms": 8.414121, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.435328, + "median": 0.064672, + "min": 0.05392, + "p90": 0.08736, + "sample_count": 2661 + }, + "host_enqueue_ms": { + "max": 1.195425, + "median": 0.156512, + "min": 0.128448, + "p90": 0.23408, + "sample_count": 2661 + }, + "sample_count": 2661, + "synchronized_e2e_ms": { + "max": 5.826406, + "median": 0.175488, + "min": 0.146528, + "p90": 0.2592, + "sample_count": 2661 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 7319, + "candidate_precomputed_gpu_span_ms": 0.013856, + "candidate_precomputed_host_enqueue_ms": 0.043296, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.013856, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.043296, + "candidate_precomputed_synchronized_e2e_ms": 0.05536, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.013856 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.036224, + "synchronized_e2e_ms": 0.052672 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.013856 + }, + "host_enqueue_ms": { + "median": 0.043296 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.013856 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 7319, + "submission_ms": { + "median": 0.043296 + }, + "synchronized_e2e_ms": { + "median": 0.05536 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7c72be90", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7c72bec0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.729792147806005, + "submission": 1.1101256467110125, + "synchronized_e2e": 1.5803468208092486 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 2661, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.037824, + "candidate_public_raw_host_enqueue_ms": 0.048064, + "candidate_public_raw_inter_kernel_gap_ms": 0.00016, + "candidate_public_raw_kernel_sum_ms": 0.037664, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.048064, + "candidate_public_raw_synchronized_e2e_ms": 0.087488, + "candidate_public_raw_tflops_from_gpu_span": 130.40665989847713, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.037664 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.245568, + "synchronized_e2e_ms": 0.275584 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.037824 + }, + "host_enqueue_ms": { + "median": 0.048064 + }, + "inter_kernel_gap_ms": { + "median": 0.00016 + }, + "kernel_sum_ms": { + "median": 0.037664 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2661, + "submission_ms": { + "median": 0.048064 + }, + "synchronized_e2e_ms": { + "median": 0.087488 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.190499, + "after_init_synchronized_e2e_ms_per_call": 3.216291, + "including_init_host_enqueue_ms_per_call": 38.064583, + "including_init_synchronized_e2e_ms_per_call": 38.159366999999996, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.3623075, + "after_init_synchronized_e2e_ms_per_call": 0.40036829999999995, + "including_init_host_enqueue_ms_per_call": 3.8497158999999996, + "including_init_synchronized_e2e_ms_per_call": 3.8946758999999993, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07948835, + "after_init_synchronized_e2e_ms_per_call": 0.11877602999999999, + "including_init_host_enqueue_ms_per_call": 0.42822919, + "including_init_synchronized_e2e_ms_per_call": 0.46820679, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.05120643500000001, + "after_init_synchronized_e2e_ms_per_call": 0.090616803, + "including_init_host_enqueue_ms_per_call": 0.08608051900000001, + "including_init_synchronized_e2e_ms_per_call": 0.12555987899999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.190499, + "synchronized_e2e_ms": 3.216291, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.03872, + "median": 0.037824, + "min": 0.037152, + "p90": 0.038144, + "sample_count": 2661 + }, + "host_enqueue_ms": { + "max": 25.731354, + "median": 0.048064, + "min": 0.037664, + "p90": 0.075776, + "sample_count": 2661 + }, + "sample_count": 2661, + "synchronized_e2e_ms": { + "max": 30.907072, + "median": 0.087488, + "min": 0.078368, + "p90": 0.1112, + "sample_count": 2661 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.07104, + "submission_ms": 0.07104, + "synchronized_e2e_ms": 0.093376 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.04544, + "submission_ms": 0.04544, + "synchronized_e2e_ms": 0.068928 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.220545, + "submission_ms": 0.220545, + "synchronized_e2e_ms": 0.239457 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.384745, + "submission_ms": 8.384745, + "synchronized_e2e_ms": 8.414121 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.066049, + "submission_ms": 0.066049, + "synchronized_e2e_ms": 0.083585 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.036224, + "submission_ms": 0.036224, + "synchronized_e2e_ms": 0.052672 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.345761, + "submission_ms": 1.345761, + "synchronized_e2e_ms": 1.367649 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.118753, + "submission_ms": 1.118753, + "synchronized_e2e_ms": 1.140033 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.245568, + "submission_ms": 0.245568, + "synchronized_e2e_ms": 0.275584 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.190499, + "submission_ms": 3.190499, + "synchronized_e2e_ms": 3.216291 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.333376, + "evolution_kernel_ms": 0.152256, + "evolution_speedup": 2.1896, + "evolution_tflops": 32.3961, + "expected_route": "aligned_weave_v10_fallback", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_c44f_d128_forced_fallback_b6_n6272_k512_d128", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 19960, + "measurement_schedule_sha256": "fa3fbe54ac1fc9b99b9a21916a6d95a497c2766a956cb4e9ea3a1d0353f74e52", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 7319, + "public_pair_count": 2661, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 7319, + "baseline_public_raw": 2661, + "candidate_precomputed": 7319, + "candidate_public_raw": 2661 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 3994 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:adjacent_c44f_d128_forced_fallback_b6_n6272_k512_d128", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.3163972286374133, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.0058522311631313, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.6160944392158543, + "including_init_synchronized_e2e_speedup": 12.905235430136985, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.496079984354406, + "including_init_synchronized_e2e_speedup": 12.684881096267858, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.1710974007129216, + "including_init_synchronized_e2e_speedup": 10.888969102733432, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.0275117518767463, + "including_init_synchronized_e2e_speedup": 5.318324072293827, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.7098138747884937, + "hot_synchronized_e2e_speedup": 2.0058522311631313, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 4412801, + "selected_route": "aligned_weave_v10_fallback", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_c44f_d128_forced_fallback_b6_n6272_k512_d128", + "source": "forced_fallback", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 224, + "K": 512, + "N": 4480, + "baseline_07cf_adapter_bench_iters": 2618, + "baseline_07cf_adapter_gpu_span_ms": 0.064224, + "baseline_07cf_adapter_host_enqueue_ms": 0.15832000000000002, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.021712000000000002, + "baseline_07cf_adapter_kernel_sum_ms": 0.042528, + "baseline_07cf_adapter_submission_ms": 0.15832000000000002, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.1772645, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.042528 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.216992, + "synchronized_e2e_ms": 0.236928 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.064224 + }, + "host_enqueue_ms": { + "median": 0.15832000000000002 + }, + "inter_kernel_gap_ms": { + "median": 0.021712000000000002 + }, + "kernel_sum_ms": { + "median": 0.042528 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2618, + "submission_ms": { + "median": 0.15832000000000002 + }, + "synchronized_e2e_ms": { + "median": 0.1772645 + } + }, + "baseline_07cf_precomputed_bench_iters": 5204, + "baseline_07cf_precomputed_gpu_span_ms": 0.019264, + "baseline_07cf_precomputed_host_enqueue_ms": 0.044256, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.019264, + "baseline_07cf_precomputed_submission_ms": 0.044256, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.068992, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.019264 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.047872, + "synchronized_e2e_ms": 0.070144 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.019264 + }, + "host_enqueue_ms": { + "median": 0.044256 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.019264 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5204, + "submission_ms": { + "median": 0.044256 + }, + "synchronized_e2e_ms": { + "median": 0.068992 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 3.333887043189369, + "submission": 3.5773680404916854, + "synchronized_e2e": 2.5693486201298703 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.342759, + "after_init_synchronized_e2e_ms_per_call": 7.369959, + "including_init_host_enqueue_ms_per_call": 42.894027, + "including_init_synchronized_e2e_ms_per_call": 43.001196, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8767639, + "after_init_synchronized_e2e_ms_per_call": 0.8965339499999999, + "including_init_host_enqueue_ms_per_call": 4.4318907, + "including_init_synchronized_e2e_ms_per_call": 4.45965765, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23016439000000002, + "after_init_synchronized_e2e_ms_per_call": 0.249191445, + "including_init_host_enqueue_ms_per_call": 0.58567707, + "including_init_synchronized_e2e_ms_per_call": 0.6055038150000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.16550443900000003, + "after_init_synchronized_e2e_ms_per_call": 0.18445719449999998, + "including_init_host_enqueue_ms_per_call": 0.201055707, + "including_init_synchronized_e2e_ms_per_call": 0.2200884315, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.342759, + "synchronized_e2e_ms": 7.369959, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.410656, + "median": 0.064224, + "min": 0.05392, + "p90": 0.08662400000000005, + "sample_count": 2618 + }, + "host_enqueue_ms": { + "max": 32.203746, + "median": 0.15832000000000002, + "min": 0.129632, + "p90": 0.2048928, + "sample_count": 2618 + }, + "sample_count": 2618, + "synchronized_e2e_ms": { + "max": 32.23613, + "median": 0.1772645, + "min": 0.145216, + "p90": 0.2278823, + "sample_count": 2618 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 5204, + "candidate_precomputed_gpu_span_ms": 0.01648, + "candidate_precomputed_host_enqueue_ms": 0.055008, + "candidate_precomputed_inter_kernel_gap_ms": 0.002272, + "candidate_precomputed_kernel_sum_ms": 0.014208, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.055008, + "candidate_precomputed_synchronized_e2e_ms": 0.066784, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.014208 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.04816, + "synchronized_e2e_ms": 0.064928 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.01648 + }, + "host_enqueue_ms": { + "median": 0.055008 + }, + "inter_kernel_gap_ms": { + "median": 0.002272 + }, + "kernel_sum_ms": { + "median": 0.014208 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5204, + "submission_ms": { + "median": 0.055008 + }, + "synchronized_e2e_ms": { + "median": 0.066784 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01232db0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01233dd0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.3242718446601938, + "submission": 0.8574752763234439, + "synchronized_e2e": 1.2889314805941545 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2618, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.038304, + "candidate_public_raw_host_enqueue_ms": 0.047168, + "candidate_public_raw_inter_kernel_gap_ms": 3.2e-05, + "candidate_public_raw_kernel_sum_ms": 0.0384, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.047168, + "candidate_public_raw_synchronized_e2e_ms": 0.08608, + "candidate_public_raw_tflops_from_gpu_span": 107.31040935672516, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.038272 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.05856, + "synchronized_e2e_ms": 0.083008 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.038304 + }, + "host_enqueue_ms": { + "median": 0.047168 + }, + "inter_kernel_gap_ms": { + "median": 3.2e-05 + }, + "kernel_sum_ms": { + "median": 0.0384 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2618, + "submission_ms": { + "median": 0.047168 + }, + "synchronized_e2e_ms": { + "median": 0.08608 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.425315, + "after_init_synchronized_e2e_ms_per_call": 2.453859, + "including_init_host_enqueue_ms_per_call": 38.263335, + "including_init_synchronized_e2e_ms_per_call": 486.554072, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2849827, + "after_init_synchronized_e2e_ms_per_call": 0.3228579, + "including_init_host_enqueue_ms_per_call": 3.8687847, + "including_init_synchronized_e2e_ms_per_call": 48.7328792, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07094947, + "after_init_synchronized_e2e_ms_per_call": 0.10975779, + "including_init_host_enqueue_ms_per_call": 0.42932967, + "including_init_synchronized_e2e_ms_per_call": 4.95075992, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.049546147, + "after_init_synchronized_e2e_ms_per_call": 0.08844777899999999, + "including_init_host_enqueue_ms_per_call": 0.085384167, + "including_init_synchronized_e2e_ms_per_call": 0.5725479920000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.425315, + "synchronized_e2e_ms": 2.453859, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.03904, + "median": 0.038304, + "min": 0.03776, + "p90": 0.038496, + "sample_count": 2618 + }, + "host_enqueue_ms": { + "max": 0.413376, + "median": 0.047168, + "min": 0.038208, + "p90": 0.06821120000000001, + "sample_count": 2618 + }, + "sample_count": 2618, + "synchronized_e2e_ms": { + "max": 43.604717, + "median": 0.08608, + "min": 0.078464, + "p90": 0.10486150000000001, + "sample_count": 2618 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.071712, + "submission_ms": 0.071712, + "synchronized_e2e_ms": 0.091616 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.047872, + "submission_ms": 0.047872, + "synchronized_e2e_ms": 0.070144 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.216992, + "submission_ms": 0.216992, + "synchronized_e2e_ms": 0.236928 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.342759, + "submission_ms": 7.342759, + "synchronized_e2e_ms": 7.369959 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.071872, + "submission_ms": 0.071872, + "synchronized_e2e_ms": 0.091968 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.04816, + "submission_ms": 0.04816, + "synchronized_e2e_ms": 0.064928 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.262977, + "submission_ms": 1.262977, + "synchronized_e2e_ms": 1.289025 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.19581, + "submission_ms": 1.19581, + "synchronized_e2e_ms": 1.23709 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.05856, + "submission_ms": 0.05856, + "synchronized_e2e_ms": 0.083008 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.425315, + "submission_ms": 2.425315, + "synchronized_e2e_ms": 2.453859 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.402944, + "evolution_kernel_ms": 0.173008, + "evolution_speedup": 2.329, + "evolution_tflops": 23.7585, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_c44f_d224_overlap_b4_n4480_k512_d224", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 15644, + "measurement_schedule_sha256": "c03c0f10d5015cd80e41e73424e6e918df33d45e6228a49969b50e9ddf38aab3", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 5204, + "public_pair_count": 2618, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 5204, + "baseline_public_raw": 2618, + "candidate_precomputed": 5204, + "candidate_public_raw": 2618 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 3130 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:adjacent_c44f_d224_overlap_b4_n4480_k512_d224", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.1689320388349513, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.0592994888475835, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.0034158441866463, + "including_init_synchronized_e2e_speedup": 0.08837906920240511, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.776868554246311, + "including_init_synchronized_e2e_speedup": 0.09151229566587972, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.270375934136429, + "including_init_synchronized_e2e_speedup": 0.12230522683071249, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.085492666808513, + "including_init_synchronized_e2e_speedup": 0.3844017175419593, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.6766917293233083, + "hot_synchronized_e2e_speedup": 2.0592994888475835, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 4422401, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_c44f_d224_overlap_b4_n4480_k512_d224", + "source": "guard_overlap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 288, + "K": 4096, + "N": 768, + "baseline_07cf_adapter_bench_iters": 2327, + "baseline_07cf_adapter_gpu_span_ms": 0.147168, + "baseline_07cf_adapter_host_enqueue_ms": 0.152896, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.030272, + "baseline_07cf_adapter_kernel_sum_ms": 0.116865, + "baseline_07cf_adapter_submission_ms": 0.152896, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.250528, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.116865 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.19168, + "synchronized_e2e_ms": 0.285504 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.147168 + }, + "host_enqueue_ms": { + "median": 0.152896 + }, + "inter_kernel_gap_ms": { + "median": 0.030272 + }, + "kernel_sum_ms": { + "median": 0.116865 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2327, + "submission_ms": { + "median": 0.152896 + }, + "synchronized_e2e_ms": { + "median": 0.250528 + } + }, + "baseline_07cf_precomputed_bench_iters": 3313, + "baseline_07cf_precomputed_gpu_span_ms": 0.121632, + "baseline_07cf_precomputed_host_enqueue_ms": 0.044224, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.121632, + "baseline_07cf_precomputed_submission_ms": 0.044224, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.17136, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.121632 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.042912, + "synchronized_e2e_ms": 0.14016 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.121632 + }, + "host_enqueue_ms": { + "median": 0.044224 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.121632 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3313, + "submission_ms": { + "median": 0.044224 + }, + "synchronized_e2e_ms": { + "median": 0.17136 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.2099447513812154, + "submission": 3.4573082489146167, + "synchronized_e2e": 1.4619981325863676 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.745224, + "after_init_synchronized_e2e_ms_per_call": 7.82532, + "including_init_host_enqueue_ms_per_call": 44.679662, + "including_init_synchronized_e2e_ms_per_call": 462.936639, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9121288, + "after_init_synchronized_e2e_ms_per_call": 1.0080072, + "including_init_host_enqueue_ms_per_call": 4.6055726, + "including_init_synchronized_e2e_ms_per_call": 46.519139100000004, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22881928, + "after_init_synchronized_e2e_ms_per_call": 0.32627592, + "including_init_host_enqueue_ms_per_call": 0.59816366, + "including_init_synchronized_e2e_ms_per_call": 4.87738911, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.160488328, + "after_init_synchronized_e2e_ms_per_call": 0.25810279199999997, + "including_init_host_enqueue_ms_per_call": 0.19742276600000003, + "including_init_synchronized_e2e_ms_per_call": 0.713214111, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.745224, + "synchronized_e2e_ms": 7.82532, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.761442, + "median": 0.147168, + "min": 0.137824, + "p90": 0.1682688, + "sample_count": 2327 + }, + "host_enqueue_ms": { + "max": 0.835361, + "median": 0.152896, + "min": 0.128288, + "p90": 0.19672380000000006, + "sample_count": 2327 + }, + "sample_count": 2327, + "synchronized_e2e_ms": { + "max": 0.895681, + "median": 0.250528, + "min": 0.229632, + "p90": 0.29011240000000005, + "sample_count": 2327 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3313, + "candidate_precomputed_gpu_span_ms": 0.030272, + "candidate_precomputed_host_enqueue_ms": 0.051648, + "candidate_precomputed_inter_kernel_gap_ms": 0.00224, + "candidate_precomputed_kernel_sum_ms": 0.028064, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.051648, + "candidate_precomputed_synchronized_e2e_ms": 0.065024, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.028064 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.037728, + "synchronized_e2e_ms": 0.052128 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.030272 + }, + "host_enqueue_ms": { + "median": 0.051648 + }, + "inter_kernel_gap_ms": { + "median": 0.00224 + }, + "kernel_sum_ms": { + "median": 0.028064 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3313, + "submission_ms": { + "median": 0.051648 + }, + "synchronized_e2e_ms": { + "median": 0.065024 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb6bcbce90", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc29535640" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.4154334038054968, + "submission": 0.9684014869888475, + "synchronized_e2e": 1.4419291338582676 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2327, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.042848, + "candidate_public_raw_host_enqueue_ms": 0.050016, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.04304, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.050016, + "candidate_public_raw_synchronized_e2e_ms": 0.09376, + "candidate_public_raw_tflops_from_gpu_span": 84.57521135175504, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.042848 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.288929, + "synchronized_e2e_ms": 0.314305 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.042848 + }, + "host_enqueue_ms": { + "median": 0.050016 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.04304 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2327, + "submission_ms": { + "median": 0.050016 + }, + "synchronized_e2e_ms": { + "median": 0.09376 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.457474, + "after_init_synchronized_e2e_ms_per_call": 2.483458, + "including_init_host_enqueue_ms_per_call": 39.81549699999999, + "including_init_synchronized_e2e_ms_per_call": 39.919081, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.29076179999999996, + "after_init_synchronized_e2e_ms_per_call": 0.33272979999999996, + "including_init_host_enqueue_ms_per_call": 4.0265641, + "including_init_synchronized_e2e_ms_per_call": 4.0762921, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07409058, + "after_init_synchronized_e2e_ms_per_call": 0.11765698000000001, + "including_init_host_enqueue_ms_per_call": 0.4476708099999999, + "including_init_synchronized_e2e_ms_per_call": 0.49201321, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.052423458, + "after_init_synchronized_e2e_ms_per_call": 0.096149698, + "including_init_host_enqueue_ms_per_call": 0.08978148099999998, + "including_init_synchronized_e2e_ms_per_call": 0.133585321, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.457474, + "synchronized_e2e_ms": 2.483458, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.043969, + "median": 0.042848, + "min": 0.041952, + "p90": 0.04336, + "sample_count": 2327 + }, + "host_enqueue_ms": { + "max": 10.060266, + "median": 0.050016, + "min": 0.04, + "p90": 0.07148800000000001, + "sample_count": 2327 + }, + "sample_count": 2327, + "synchronized_e2e_ms": { + "max": 10.09089, + "median": 0.09376, + "min": 0.085152, + "p90": 0.11329280000000003, + "sample_count": 2327 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.065824, + "submission_ms": 0.065824, + "synchronized_e2e_ms": 0.161568 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.042912, + "submission_ms": 0.042912, + "synchronized_e2e_ms": 0.14016 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.19168, + "submission_ms": 0.19168, + "synchronized_e2e_ms": 0.285504 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.745224, + "submission_ms": 7.745224, + "synchronized_e2e_ms": 7.82532 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.088416, + "submission_ms": 0.088416, + "synchronized_e2e_ms": 0.10816 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.037728, + "submission_ms": 0.037728, + "synchronized_e2e_ms": 0.052128 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.316929, + "submission_ms": 1.316929, + "synchronized_e2e_ms": 1.337857 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.235426, + "submission_ms": 1.235426, + "synchronized_e2e_ms": 1.256546 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.288929, + "submission_ms": 0.288929, + "synchronized_e2e_ms": 0.314305 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.457474, + "submission_ms": 2.457474, + "synchronized_e2e_ms": 2.483458 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 1.583007, + "evolution_kernel_ms": 0.246496, + "evolution_speedup": 6.422, + "evolution_tflops": 14.7016, + "expected_route": "d288_parent_splitk_hybrid_20260629_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_c44f_d288_highk_heldout_b2_n768_k4096_d288", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 11280, + "measurement_schedule_sha256": "a6c92e2077a6224fc77e711af6732cf35b383b85398c35a8d4830d6e61d3b304", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3313, + "public_pair_count": 2327, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3313, + "baseline_public_raw": 2327, + "candidate_precomputed": 3313, + "candidate_public_raw": 2327 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2258 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:adjacent_c44f_d288_highk_heldout_b2_n768_k4096_d288", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 4.017970401691332, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.672013651877133, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.150977387175462, + "including_init_synchronized_e2e_speedup": 11.596876165561026, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 3.0295068250574495, + "including_init_synchronized_e2e_speedup": 11.41212110388262, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.773111463510282, + "including_init_synchronized_e2e_speedup": 9.91312633658759, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.684384843309648, + "including_init_synchronized_e2e_speedup": 5.339015586899701, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 3.434652725914862, + "hot_synchronized_e2e_speedup": 2.672013651877133, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 4428801, + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_c44f_d288_highk_heldout_b2_n768_k4096_d288", + "source": "heldout_neighborhood", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 352, + "K": 768, + "N": 3328, + "baseline_07cf_adapter_bench_iters": 3078, + "baseline_07cf_adapter_gpu_span_ms": 0.067104, + "baseline_07cf_adapter_host_enqueue_ms": 0.14494400000000002, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.034784, + "baseline_07cf_adapter_kernel_sum_ms": 0.032288, + "baseline_07cf_adapter_submission_ms": 0.14494400000000002, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.166848, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.032288 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.203168, + "synchronized_e2e_ms": 0.2224 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.067104 + }, + "host_enqueue_ms": { + "median": 0.14494400000000002 + }, + "inter_kernel_gap_ms": { + "median": 0.034784 + }, + "kernel_sum_ms": { + "median": 0.032288 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3078, + "submission_ms": { + "median": 0.14494400000000002 + }, + "synchronized_e2e_ms": { + "median": 0.166848 + } + }, + "baseline_07cf_precomputed_bench_iters": 3677, + "baseline_07cf_precomputed_gpu_span_ms": 0.027136, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042272, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.027136, + "baseline_07cf_precomputed_submission_ms": 0.042272, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.076, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.027136 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.041184, + "synchronized_e2e_ms": 0.067872 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.027136 + }, + "host_enqueue_ms": { + "median": 0.042272 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.027136 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3677, + "submission_ms": { + "median": 0.042272 + }, + "synchronized_e2e_ms": { + "median": 0.076 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.4728773584905657, + "submission": 3.42884178652536, + "synchronized_e2e": 2.1953684210526316 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.90548, + "after_init_synchronized_e2e_ms_per_call": 7.929704, + "including_init_host_enqueue_ms_per_call": 42.133387, + "including_init_synchronized_e2e_ms_per_call": 42.275147000000004, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9209976, + "after_init_synchronized_e2e_ms_per_call": 0.9431336, + "including_init_host_enqueue_ms_per_call": 4.3437883, + "including_init_synchronized_e2e_ms_per_call": 4.3776779, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22254936, + "after_init_synchronized_e2e_ms_per_call": 0.24447656, + "including_init_host_enqueue_ms_per_call": 0.56482843, + "including_init_synchronized_e2e_ms_per_call": 0.58793099, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15270453600000003, + "after_init_synchronized_e2e_ms_per_call": 0.17461085599999998, + "including_init_host_enqueue_ms_per_call": 0.186932443, + "including_init_synchronized_e2e_ms_per_call": 0.208956299, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.90548, + "synchronized_e2e_ms": 7.929704, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.231302, + "median": 0.067104, + "min": 0.057888, + "p90": 0.08572510000000001, + "sample_count": 3078 + }, + "host_enqueue_ms": { + "max": 50.032851, + "median": 0.14494400000000002, + "min": 0.120096, + "p90": 0.2159203, + "sample_count": 3078 + }, + "sample_count": 3078, + "synchronized_e2e_ms": { + "max": 50.169587, + "median": 0.166848, + "min": 0.143712, + "p90": 0.24114560000000004, + "sample_count": 3078 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3677, + "candidate_precomputed_gpu_span_ms": 0.027776, + "candidate_precomputed_host_enqueue_ms": 0.053952, + "candidate_precomputed_inter_kernel_gap_ms": 0.004192, + "candidate_precomputed_kernel_sum_ms": 0.023616, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.053952, + "candidate_precomputed_synchronized_e2e_ms": 0.066272, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.023616 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.037952, + "synchronized_e2e_ms": 0.053792 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.027776 + }, + "host_enqueue_ms": { + "median": 0.053952 + }, + "inter_kernel_gap_ms": { + "median": 0.004192 + }, + "kernel_sum_ms": { + "median": 0.023616 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3677, + "submission_ms": { + "median": 0.053952 + }, + "synchronized_e2e_ms": { + "median": 0.066272 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc03a5c770", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc03a5f4d0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.1716589861751152, + "submission": 0.8131672597864769, + "synchronized_e2e": 1.183003380009657 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 3078, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.032544, + "candidate_public_raw_host_enqueue_ms": 0.043872, + "candidate_public_raw_inter_kernel_gap_ms": 9.6e-05, + "candidate_public_raw_kernel_sum_ms": 0.032448, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.043872, + "candidate_public_raw_synchronized_e2e_ms": 0.0784, + "candidate_public_raw_tflops_from_gpu_span": 55.2899587020649, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.032448 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.071649, + "synchronized_e2e_ms": 0.096705 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.032544 + }, + "host_enqueue_ms": { + "median": 0.043872 + }, + "inter_kernel_gap_ms": { + "median": 9.6e-05 + }, + "kernel_sum_ms": { + "median": 0.032448 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3078, + "submission_ms": { + "median": 0.043872 + }, + "synchronized_e2e_ms": { + "median": 0.0784 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.339842, + "after_init_synchronized_e2e_ms_per_call": 2.363202, + "including_init_host_enqueue_ms_per_call": 36.917381999999996, + "including_init_synchronized_e2e_ms_per_call": 450.450994, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.273469, + "after_init_synchronized_e2e_ms_per_call": 0.3068802, + "including_init_host_enqueue_ms_per_call": 3.731223, + "including_init_synchronized_e2e_ms_per_call": 45.1156594, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06683170000000001, + "after_init_synchronized_e2e_ms_per_call": 0.10124802, + "including_init_host_enqueue_ms_per_call": 0.41260709999999995, + "including_init_synchronized_e2e_ms_per_call": 4.58212594, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.046167969999999996, + "after_init_synchronized_e2e_ms_per_call": 0.080684802, + "including_init_host_enqueue_ms_per_call": 0.08074550999999999, + "including_init_synchronized_e2e_ms_per_call": 0.528772594, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.339842, + "synchronized_e2e_ms": 2.363202, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.03312, + "median": 0.032544, + "min": 0.032128, + "p90": 0.032736, + "sample_count": 3078 + }, + "host_enqueue_ms": { + "max": 50.494324, + "median": 0.043872, + "min": 0.03504, + "p90": 0.06629190000000001, + "sample_count": 3078 + }, + "sample_count": 3078, + "synchronized_e2e_ms": { + "max": 50.710356, + "median": 0.0784, + "min": 0.07072, + "p90": 0.096608, + "sample_count": 3078 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.061664, + "submission_ms": 0.061664, + "synchronized_e2e_ms": 0.086784 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.041184, + "submission_ms": 0.041184, + "synchronized_e2e_ms": 0.067872 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.203168, + "submission_ms": 0.203168, + "synchronized_e2e_ms": 0.2224 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.90548, + "submission_ms": 7.90548, + "synchronized_e2e_ms": 7.929704 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.069248, + "submission_ms": 0.069248, + "synchronized_e2e_ms": 0.086048 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.037952, + "submission_ms": 0.037952, + "synchronized_e2e_ms": 0.053792 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.176897, + "submission_ms": 1.176897, + "synchronized_e2e_ms": 1.200065 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.046177, + "submission_ms": 1.046177, + "synchronized_e2e_ms": 1.068097 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.071649, + "submission_ms": 0.071649, + "synchronized_e2e_ms": 0.096705 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.339842, + "submission_ms": 2.339842, + "synchronized_e2e_ms": 2.363202 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.4864, + "evolution_kernel_ms": 0.187632, + "evolution_speedup": 2.5923, + "evolution_tflops": 9.5898, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_c44f_d352_random_b1_n3328_k768_d352", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 13510, + "measurement_schedule_sha256": "2144e831fa1a005adb709b4090ff8b30a89875d7b9f3e4a735882793b02dcd57", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3677, + "public_pair_count": 3078, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3677, + "baseline_public_raw": 3078, + "candidate_precomputed": 3677, + "candidate_public_raw": 3078 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2704 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:adjacent_c44f_d352_random_b1_n3328_k768_d352", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.9769585253456222, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.1281632653061227, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.3554914053051754, + "including_init_synchronized_e2e_speedup": 0.09385071309222154, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 3.0732957030137493, + "including_init_synchronized_e2e_speedup": 0.09703233773415712, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.4146305280834137, + "including_init_synchronized_e2e_speedup": 0.12830965313886594, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.1641108569616367, + "including_init_synchronized_e2e_speedup": 0.39517233187013473, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.0619469026548676, + "hot_synchronized_e2e_speedup": 2.1281632653061227, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 4435201, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_c44f_d352_random_b1_n3328_k768_d352", + "source": "random_legal", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 3, + "D": 416, + "K": 8192, + "N": 512, + "baseline_07cf_adapter_bench_iters": 1176, + "baseline_07cf_adapter_gpu_span_ms": 0.25984, + "baseline_07cf_adapter_host_enqueue_ms": 0.1521925, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.005776, + "baseline_07cf_adapter_kernel_sum_ms": 0.254048, + "baseline_07cf_adapter_submission_ms": 0.1521925, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.36419250000000003, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.254048 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.205664, + "synchronized_e2e_ms": 0.410528 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.25984 + }, + "host_enqueue_ms": { + "median": 0.1521925 + }, + "inter_kernel_gap_ms": { + "median": 0.005776 + }, + "kernel_sum_ms": { + "median": 0.254048 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1176, + "submission_ms": { + "median": 0.1521925 + }, + "synchronized_e2e_ms": { + "median": 0.36419250000000003 + } + }, + "baseline_07cf_precomputed_bench_iters": 1860, + "baseline_07cf_precomputed_gpu_span_ms": 0.2563205, + "baseline_07cf_precomputed_host_enqueue_ms": 0.044, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.2563205, + "baseline_07cf_precomputed_submission_ms": 0.044, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.308129, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.2563205 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.040864, + "synchronized_e2e_ms": 0.276064 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.2563205 + }, + "host_enqueue_ms": { + "median": 0.044 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.2563205 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1860, + "submission_ms": { + "median": 0.044 + }, + "synchronized_e2e_ms": { + "median": 0.308129 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.013730856486313, + "submission": 3.458920454545455, + "synchronized_e2e": 1.181948145095074 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.829128, + "after_init_synchronized_e2e_ms_per_call": 8.019145, + "including_init_host_enqueue_ms_per_call": 42.278572, + "including_init_synchronized_e2e_ms_per_call": 492.060639, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.91988605, + "after_init_synchronized_e2e_ms_per_call": 1.12968775, + "including_init_host_enqueue_ms_per_call": 4.3648304499999995, + "including_init_synchronized_e2e_ms_per_call": 49.53383715, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22896185500000002, + "after_init_synchronized_e2e_ms_per_call": 0.44074202500000004, + "including_init_host_enqueue_ms_per_call": 0.5734562950000001, + "including_init_synchronized_e2e_ms_per_call": 5.281156965, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15986943550000002, + "after_init_synchronized_e2e_ms_per_call": 0.3718474525, + "including_init_host_enqueue_ms_per_call": 0.1943188795, + "including_init_synchronized_e2e_ms_per_call": 0.8558889465, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.829128, + "synchronized_e2e_ms": 8.019145, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.709472, + "median": 0.25984, + "min": 0.254752, + "p90": 0.2676, + "sample_count": 1176 + }, + "host_enqueue_ms": { + "max": 0.620641, + "median": 0.1521925, + "min": 0.131841, + "p90": 0.181456, + "sample_count": 1176 + }, + "sample_count": 1176, + "synchronized_e2e_ms": { + "max": 0.818689, + "median": 0.36419250000000003, + "min": 0.345825, + "p90": 0.390673, + "sample_count": 1176 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 4, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 1860, + "candidate_precomputed_gpu_span_ms": 0.05376, + "candidate_precomputed_host_enqueue_ms": 0.068864, + "candidate_precomputed_inter_kernel_gap_ms": 0.0037920000000000002, + "candidate_precomputed_kernel_sum_ms": 0.049953, + "candidate_precomputed_launch_count": 3, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.068864, + "candidate_precomputed_synchronized_e2e_ms": 0.09121599999999999, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.049953 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.047008, + "synchronized_e2e_ms": 0.078144 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 3.0 + }, + "gpu_span_ms": { + "median": 0.05376 + }, + "host_enqueue_ms": { + "median": 0.068864 + }, + "inter_kernel_gap_ms": { + "median": 0.0037920000000000002 + }, + "kernel_sum_ms": { + "median": 0.049953 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1860, + "submission_ms": { + "median": 0.068864 + }, + "synchronized_e2e_ms": { + "median": 0.09121599999999999 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01990080", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc019935c0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.5845238095238094, + "submission": 0.7132899628252788, + "synchronized_e2e": 1.506583274864059 + }, + "candidate_public_raw_assignment_launch_count": 3, + "candidate_public_raw_bench_iters": 1176, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.085184, + "candidate_public_raw_host_enqueue_ms": 0.04912, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.085536, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.04912, + "candidate_public_raw_synchronized_e2e_ms": 0.1374245, + "candidate_public_raw_tflops_from_gpu_span": 122.89846431254696, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.085184 + }, + "activity_count": { + "median": 4.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.070688, + "synchronized_e2e_ms": 0.144672 + }, + "correlated_kernel_activity_count": { + "median": 4.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.085184 + }, + "host_enqueue_ms": { + "median": 0.04912 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.085536 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1176, + "submission_ms": { + "median": 0.04912 + }, + "synchronized_e2e_ms": { + "median": 0.1374245 + } + }, + "candidate_public_raw_total_launch_count": 4, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.477219, + "after_init_synchronized_e2e_ms_per_call": 3.533891, + "including_init_host_enqueue_ms_per_call": 38.351302999999994, + "including_init_synchronized_e2e_ms_per_call": 38.476966999999995, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.39192989999999994, + "after_init_synchronized_e2e_ms_per_call": 0.47707115, + "including_init_host_enqueue_ms_per_call": 3.8793382999999992, + "including_init_synchronized_e2e_ms_per_call": 3.9713787499999995, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.08340098999999998, + "after_init_synchronized_e2e_ms_per_call": 0.171389165, + "including_init_host_enqueue_ms_per_call": 0.4321418299999999, + "including_init_synchronized_e2e_ms_per_call": 0.520819925, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.052548098999999994, + "after_init_synchronized_e2e_ms_per_call": 0.14082096650000003, + "including_init_host_enqueue_ms_per_call": 0.08742218299999999, + "including_init_synchronized_e2e_ms_per_call": 0.17576404250000002, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.477219, + "synchronized_e2e_ms": 3.533891, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.086593, + "median": 0.085184, + "min": 0.084448, + "p90": 0.085632, + "sample_count": 1176 + }, + "host_enqueue_ms": { + "max": 0.102304, + "median": 0.04912, + "min": 0.040704, + "p90": 0.0591685, + "sample_count": 1176 + }, + "sample_count": 1176, + "synchronized_e2e_ms": { + "max": 0.197856, + "median": 0.1374245, + "min": 0.129504, + "p90": 0.1459365, + "sample_count": 1176 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.06272, + "submission_ms": 0.06272, + "synchronized_e2e_ms": 0.320897 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.040864, + "submission_ms": 0.040864, + "synchronized_e2e_ms": 0.276064 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.205664, + "submission_ms": 0.205664, + "synchronized_e2e_ms": 0.410528 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.829128, + "submission_ms": 7.829128, + "synchronized_e2e_ms": 8.019145 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.08576, + "submission_ms": 0.08576, + "synchronized_e2e_ms": 0.105312 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.047008, + "submission_ms": 0.047008, + "synchronized_e2e_ms": 0.078144 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.520738, + "submission_ms": 1.520738, + "synchronized_e2e_ms": 1.54429 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.278401, + "submission_ms": 1.278401, + "synchronized_e2e_ms": 1.299745 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.070688, + "submission_ms": 0.070688, + "synchronized_e2e_ms": 0.144672 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.477219, + "submission_ms": 3.477219, + "synchronized_e2e_ms": 3.533891 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 3.822732, + "evolution_kernel_ms": 0.303423, + "evolution_speedup": 12.5987, + "evolution_tflops": 34.5029, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_c44f_d416_highk_request_b3_n512_k8192_d416", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 6072, + "measurement_schedule_sha256": "31d82b6bb79c54ac896d3608205a5c625f8e06c015f2b4fadfae0248782b3c4e", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 1860, + "public_pair_count": 1176, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 1860, + "baseline_public_raw": 1176, + "candidate_precomputed": 1860, + "candidate_public_raw": 1176 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 1216 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:adjacent_c44f_d416_highk_request_b3_n512_k8192_d416", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 4.767866443452381, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.6501278884041786, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.269211189592435, + "including_init_synchronized_e2e_speedup": 12.788446630941571, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.3679649251479575, + "including_init_synchronized_e2e_speedup": 12.472705392302359, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.571586278514164, + "including_init_synchronized_e2e_speedup": 10.140082419081798, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.6405688140196077, + "including_init_synchronized_e2e_speedup": 4.86953380410558, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 3.050338091660406, + "hot_synchronized_e2e_speedup": 2.6501278884041786, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 4441601, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_c44f_d416_highk_request_b3_n512_k8192_d416", + "source": "request_specific", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 5, + "D": 480, + "K": 512, + "N": 2048, + "baseline_07cf_adapter_bench_iters": 2189, + "baseline_07cf_adapter_gpu_span_ms": 0.078528, + "baseline_07cf_adapter_host_enqueue_ms": 0.152128, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.023744, + "baseline_07cf_adapter_kernel_sum_ms": 0.054496, + "baseline_07cf_adapter_submission_ms": 0.152128, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.180608, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.054496 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.229952, + "synchronized_e2e_ms": 0.254208 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.078528 + }, + "host_enqueue_ms": { + "median": 0.152128 + }, + "inter_kernel_gap_ms": { + "median": 0.023744 + }, + "kernel_sum_ms": { + "median": 0.054496 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2189, + "submission_ms": { + "median": 0.152128 + }, + "synchronized_e2e_ms": { + "median": 0.180608 + } + }, + "baseline_07cf_precomputed_bench_iters": 3442, + "baseline_07cf_precomputed_gpu_span_ms": 0.038497, + "baseline_07cf_precomputed_host_enqueue_ms": 0.043712, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.038497, + "baseline_07cf_precomputed_submission_ms": 0.043712, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.087808, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.038497 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.051424, + "synchronized_e2e_ms": 0.088064 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.038497 + }, + "host_enqueue_ms": { + "median": 0.043712 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.038497 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3442, + "submission_ms": { + "median": 0.043712 + }, + "synchronized_e2e_ms": { + "median": 0.087808 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.039847260825519, + "submission": 3.4802342606149343, + "synchronized_e2e": 2.056851311953353 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.319848, + "after_init_synchronized_e2e_ms_per_call": 7.34884, + "including_init_host_enqueue_ms_per_call": 42.871116, + "including_init_synchronized_e2e_ms_per_call": 42.980077, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8689, + "after_init_synchronized_e2e_ms_per_call": 0.8974312, + "including_init_host_enqueue_ms_per_call": 4.4240268, + "including_init_synchronized_e2e_ms_per_call": 4.4605549, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22380520000000004, + "after_init_synchronized_e2e_ms_per_call": 0.25229031999999996, + "including_init_host_enqueue_ms_per_call": 0.5793178800000001, + "including_init_synchronized_e2e_ms_per_call": 0.60860269, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15929572000000003, + "after_init_synchronized_e2e_ms_per_call": 0.187776232, + "including_init_host_enqueue_ms_per_call": 0.194846988, + "including_init_synchronized_e2e_ms_per_call": 0.223407469, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.319848, + "synchronized_e2e_ms": 7.34884, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.572576, + "median": 0.078528, + "min": 0.068704, + "p90": 0.10144, + "sample_count": 2189 + }, + "host_enqueue_ms": { + "max": 23.345337, + "median": 0.152128, + "min": 0.125984, + "p90": 0.19791440000000005, + "sample_count": 2189 + }, + "sample_count": 2189, + "synchronized_e2e_ms": { + "max": 41.600651, + "median": 0.180608, + "min": 0.156864, + "p90": 0.22432720000000003, + "sample_count": 2189 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3442, + "candidate_precomputed_gpu_span_ms": 0.027104, + "candidate_precomputed_host_enqueue_ms": 0.053888, + "candidate_precomputed_inter_kernel_gap_ms": 0.001792, + "candidate_precomputed_kernel_sum_ms": 0.025312, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.053888, + "candidate_precomputed_synchronized_e2e_ms": 0.06592, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.025312 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.047648, + "synchronized_e2e_ms": 0.062304 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.027104 + }, + "host_enqueue_ms": { + "median": 0.053888 + }, + "inter_kernel_gap_ms": { + "median": 0.001792 + }, + "kernel_sum_ms": { + "median": 0.025312 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3442, + "submission_ms": { + "median": 0.053888 + }, + "synchronized_e2e_ms": { + "median": 0.06592 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04666510", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc046674a0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.6824085005903189, + "submission": 0.8485748218527316, + "synchronized_e2e": 1.3946601941747572 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2189, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.0456, + "candidate_public_raw_host_enqueue_ms": 0.045728, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.045536, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.045728, + "candidate_public_raw_synchronized_e2e_ms": 0.091936, + "candidate_public_raw_tflops_from_gpu_span": 110.37642105263157, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.045536 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.088064, + "synchronized_e2e_ms": 0.11232 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.0456 + }, + "host_enqueue_ms": { + "median": 0.045728 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.045536 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2189, + "submission_ms": { + "median": 0.045728 + }, + "synchronized_e2e_ms": { + "median": 0.091936 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.429698, + "after_init_synchronized_e2e_ms_per_call": 2.457442, + "including_init_host_enqueue_ms_per_call": 38.267718, + "including_init_synchronized_e2e_ms_per_call": 486.557655, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.284125, + "after_init_synchronized_e2e_ms_per_call": 0.3284866, + "including_init_host_enqueue_ms_per_call": 3.8679270000000003, + "including_init_synchronized_e2e_ms_per_call": 48.7385079, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.0695677, + "after_init_synchronized_e2e_ms_per_call": 0.11559106, + "including_init_host_enqueue_ms_per_call": 0.4279479, + "including_init_synchronized_e2e_ms_per_call": 4.9565931899999995, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.04811197, + "after_init_synchronized_e2e_ms_per_call": 0.09430150600000001, + "including_init_host_enqueue_ms_per_call": 0.08394999, + "including_init_synchronized_e2e_ms_per_call": 0.5784017189999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.429698, + "synchronized_e2e_ms": 2.457442, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.046624, + "median": 0.0456, + "min": 0.045056, + "p90": 0.046048, + "sample_count": 2189 + }, + "host_enqueue_ms": { + "max": 15.8948, + "median": 0.045728, + "min": 0.034592, + "p90": 0.065888, + "sample_count": 2189 + }, + "sample_count": 2189, + "synchronized_e2e_ms": { + "max": 21.259702, + "median": 0.091936, + "min": 0.08192, + "p90": 0.10997119999999999, + "sample_count": 2189 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.071264, + "submission_ms": 0.071264, + "synchronized_e2e_ms": 0.10816 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.051424, + "submission_ms": 0.051424, + "synchronized_e2e_ms": 0.088064 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.229952, + "submission_ms": 0.229952, + "synchronized_e2e_ms": 0.254208 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.319848, + "submission_ms": 7.319848, + "synchronized_e2e_ms": 7.34884 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.075872, + "submission_ms": 0.075872, + "synchronized_e2e_ms": 0.09248 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.047648, + "submission_ms": 0.047648, + "synchronized_e2e_ms": 0.062304 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.303746, + "submission_ms": 1.303746, + "synchronized_e2e_ms": 1.328034 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.174913, + "submission_ms": 1.174913, + "synchronized_e2e_ms": 1.199041 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.088064, + "submission_ms": 0.088064, + "synchronized_e2e_ms": 0.11232 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.429698, + "submission_ms": 2.429698, + "synchronized_e2e_ms": 2.457442 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.501759, + "evolution_kernel_ms": 0.185136, + "evolution_speedup": 2.7102, + "evolution_tflops": 27.1863, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_c44f_d480_boundary_b5_n2048_k512_d480", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 11262, + "measurement_schedule_sha256": "555a0ebac3405c74818e462e2569c8df5e2c6472a1eb46ddaf234f5bfd04ea8a", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3442, + "public_pair_count": 2189, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3442, + "baseline_public_raw": 2189, + "candidate_precomputed": 3442, + "candidate_public_raw": 2189 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2254 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:adjacent_c44f_d480_boundary_b5_n2048_k512_d480", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.42034386068477, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.9644970414201182, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.990442907706469, + "including_init_synchronized_e2e_speedup": 0.08833501345282503, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.7320176835219456, + "including_init_synchronized_e2e_speedup": 0.09152013658588017, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.18261100815236, + "including_init_synchronized_e2e_speedup": 0.12278649198563743, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.9912325896470835, + "including_init_synchronized_e2e_speedup": 0.38624966292674523, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.7221052631578946, + "hot_synchronized_e2e_speedup": 1.9644970414201182, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 4448001, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_c44f_d480_boundary_b5_n2048_k512_d480", + "source": "guard_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 3, + "D": 48, + "K": 768, + "N": 2688, + "baseline_07cf_adapter_bench_iters": 5719, + "baseline_07cf_adapter_gpu_span_ms": 0.059264, + "baseline_07cf_adapter_host_enqueue_ms": 0.159136, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.038688, + "baseline_07cf_adapter_kernel_sum_ms": 0.020576, + "baseline_07cf_adapter_submission_ms": 0.159136, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.179936, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.020576 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.213633, + "synchronized_e2e_ms": 0.236737 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.059264 + }, + "host_enqueue_ms": { + "median": 0.159136 + }, + "inter_kernel_gap_ms": { + "median": 0.038688 + }, + "kernel_sum_ms": { + "median": 0.020576 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 5719, + "submission_ms": { + "median": 0.159136 + }, + "synchronized_e2e_ms": { + "median": 0.179936 + } + }, + "baseline_07cf_precomputed_bench_iters": 6134, + "baseline_07cf_precomputed_gpu_span_ms": 0.016224, + "baseline_07cf_precomputed_host_enqueue_ms": 0.045216, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.016224, + "baseline_07cf_precomputed_submission_ms": 0.045216, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.066624, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.016224 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.048768, + "synchronized_e2e_ms": 0.065248 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.016224 + }, + "host_enqueue_ms": { + "median": 0.045216 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.016224 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 6134, + "submission_ms": { + "median": 0.045216 + }, + "synchronized_e2e_ms": { + "median": 0.066624 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 3.6528599605522682, + "submission": 3.5194621372965322, + "synchronized_e2e": 2.7007684918347743 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.565672, + "after_init_synchronized_e2e_ms_per_call": 7.594952, + "including_init_host_enqueue_ms_per_call": 44.50011, + "including_init_synchronized_e2e_ms_per_call": 462.706271, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8997896000000001, + "after_init_synchronized_e2e_ms_per_call": 0.9214376, + "including_init_host_enqueue_ms_per_call": 4.5932334, + "including_init_synchronized_e2e_ms_per_call": 46.4325695, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23320136000000002, + "after_init_synchronized_e2e_ms_per_call": 0.25408616, + "including_init_host_enqueue_ms_per_call": 0.6025457399999999, + "including_init_synchronized_e2e_ms_per_call": 4.8051993500000005, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.16654253600000002, + "after_init_synchronized_e2e_ms_per_call": 0.187351016, + "including_init_host_enqueue_ms_per_call": 0.203476974, + "including_init_synchronized_e2e_ms_per_call": 0.642462335, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.565672, + "synchronized_e2e_ms": 7.594952, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.242662, + "median": 0.059264, + "min": 0.048096, + "p90": 0.0882944, + "sample_count": 5719 + }, + "host_enqueue_ms": { + "max": 78.589137, + "median": 0.159136, + "min": 0.126912, + "p90": 0.29115519999999956, + "sample_count": 5719 + }, + "sample_count": 5719, + "synchronized_e2e_ms": { + "max": 83.758423, + "median": 0.179936, + "min": 0.143744, + "p90": 0.3197183999999996, + "sample_count": 5719 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 6134, + "candidate_precomputed_gpu_span_ms": 0.016768, + "candidate_precomputed_host_enqueue_ms": 0.057792, + "candidate_precomputed_inter_kernel_gap_ms": 0.005792, + "candidate_precomputed_kernel_sum_ms": 0.010976, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.057792, + "candidate_precomputed_synchronized_e2e_ms": 0.070528, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.010976 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.04896, + "synchronized_e2e_ms": 0.067072 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.016768 + }, + "host_enqueue_ms": { + "median": 0.057792 + }, + "inter_kernel_gap_ms": { + "median": 0.005792 + }, + "kernel_sum_ms": { + "median": 0.010976 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 6134, + "submission_ms": { + "median": 0.057792 + }, + "synchronized_e2e_ms": { + "median": 0.070528 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc049a2690", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc049a03e0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.0534351145038165, + "submission": 0.8294573643410852, + "synchronized_e2e": 0.9450998185117969 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 5719, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.017664, + "candidate_public_raw_host_enqueue_ms": 0.047936, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.017856, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.047936, + "candidate_public_raw_synchronized_e2e_ms": 0.066656, + "candidate_public_raw_tflops_from_gpu_span": 33.658434782608694, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.017664 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.07184, + "synchronized_e2e_ms": 0.091904 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.017664 + }, + "host_enqueue_ms": { + "median": 0.047936 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.017856 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 5719, + "submission_ms": { + "median": 0.047936 + }, + "synchronized_e2e_ms": { + "median": 0.066656 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.436354, + "after_init_synchronized_e2e_ms_per_call": 2.46445, + "including_init_host_enqueue_ms_per_call": 39.794377, + "including_init_synchronized_e2e_ms_per_call": 39.900073, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.28677779999999997, + "after_init_synchronized_e2e_ms_per_call": 0.30643539999999997, + "including_init_host_enqueue_ms_per_call": 4.0225801, + "including_init_synchronized_e2e_ms_per_call": 4.0499977000000005, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07182018, + "after_init_synchronized_e2e_ms_per_call": 0.09063394000000001, + "including_init_host_enqueue_ms_per_call": 0.44540040999999997, + "including_init_synchronized_e2e_ms_per_call": 0.46499017000000004, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.050324418, + "after_init_synchronized_e2e_ms_per_call": 0.06905379400000002, + "including_init_host_enqueue_ms_per_call": 0.087682441, + "including_init_synchronized_e2e_ms_per_call": 0.106489417, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.436354, + "synchronized_e2e_ms": 2.46445, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.018464, + "median": 0.017664, + "min": 0.01728, + "p90": 0.018144, + "sample_count": 5719 + }, + "host_enqueue_ms": { + "max": 60.948095, + "median": 0.047936, + "min": 0.03536, + "p90": 0.09244799999999985, + "sample_count": 5719 + }, + "sample_count": 5719, + "synchronized_e2e_ms": { + "max": 61.070559, + "median": 0.066656, + "min": 0.05504, + "p90": 0.11671679999999998, + "sample_count": 5719 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.072384, + "submission_ms": 0.072384, + "synchronized_e2e_ms": 0.092864 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.048768, + "submission_ms": 0.048768, + "synchronized_e2e_ms": 0.065248 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.213633, + "submission_ms": 0.213633, + "synchronized_e2e_ms": 0.236737 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.565672, + "submission_ms": 7.565672, + "synchronized_e2e_ms": 7.594952 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.107584, + "submission_ms": 0.107584, + "synchronized_e2e_ms": 0.127904 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.04896, + "submission_ms": 0.04896, + "synchronized_e2e_ms": 0.067072 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.316321, + "submission_ms": 1.316321, + "synchronized_e2e_ms": 1.341793 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.686882, + "submission_ms": 1.686882, + "synchronized_e2e_ms": 1.710754 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.07184, + "submission_ms": 0.07184, + "synchronized_e2e_ms": 0.091904 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.436354, + "submission_ms": 2.436354, + "synchronized_e2e_ms": 2.46445 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.309952, + "evolution_kernel_ms": 0.174383, + "evolution_speedup": 1.7774, + "evolution_tflops": 3.4094, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_c44f_d48_midk_boundary_b3_n2688_k768_d48", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 23706, + "measurement_schedule_sha256": "4b539d6af62ef80be6dfbb2f9012dcdb9e36d30362875389f096d76d58c6b1ef", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 6134, + "public_pair_count": 5719, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 6134, + "baseline_public_raw": 5719, + "candidate_precomputed": 6134, + "candidate_public_raw": 5719 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 4742 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:adjacent_c44f_d48_midk_boundary_b3_n2688_k768_d48", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.9675572519083968, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.699471915506481, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.081804053642801, + "including_init_synchronized_e2e_speedup": 11.596627179103157, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 3.006955462717428, + "including_init_synchronized_e2e_speedup": 11.46483848620457, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.8034327979121283, + "including_init_synchronized_e2e_speedup": 10.333980501136185, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.7131169070884065, + "including_init_synchronized_e2e_speedup": 6.03310970328629, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 3.3550724637681157, + "hot_synchronized_e2e_speedup": 2.699471915506481, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 4404801, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_c44f_d48_midk_boundary_b3_n2688_k768_d48", + "source": "guard_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 64, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 3, + "D": 112, + "K": 8192, + "N": 768, + "baseline_07cf_adapter_bench_iters": 1061, + "baseline_07cf_adapter_gpu_span_ms": 0.131264, + "baseline_07cf_adapter_host_enqueue_ms": 0.141152, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.025312, + "baseline_07cf_adapter_kernel_sum_ms": 0.105952, + "baseline_07cf_adapter_submission_ms": 0.141152, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.227968, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.105952 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.184576, + "synchronized_e2e_ms": 0.264832 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.131264 + }, + "host_enqueue_ms": { + "median": 0.141152 + }, + "inter_kernel_gap_ms": { + "median": 0.025312 + }, + "kernel_sum_ms": { + "median": 0.105952 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1061, + "submission_ms": { + "median": 0.141152 + }, + "synchronized_e2e_ms": { + "median": 0.227968 + } + }, + "baseline_07cf_precomputed_bench_iters": 1290, + "baseline_07cf_precomputed_gpu_span_ms": 0.10688, + "baseline_07cf_precomputed_host_enqueue_ms": 0.041712, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.10688, + "baseline_07cf_precomputed_submission_ms": 0.041712, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.155232, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.10688 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.05872, + "synchronized_e2e_ms": 0.127648 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.10688 + }, + "host_enqueue_ms": { + "median": 0.041712 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.10688 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1290, + "submission_ms": { + "median": 0.041712 + }, + "synchronized_e2e_ms": { + "median": 0.155232 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.2281437125748502, + "submission": 3.3839662447257384, + "synchronized_e2e": 1.468563182848897 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.87876, + "after_init_synchronized_e2e_ms_per_call": 7.94308, + "including_init_host_enqueue_ms_per_call": 42.106667, + "including_init_synchronized_e2e_ms_per_call": 42.288523000000005, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9149128, + "after_init_synchronized_e2e_ms_per_call": 0.9994792, + "including_init_host_enqueue_ms_per_call": 4.3377035, + "including_init_synchronized_e2e_ms_per_call": 4.4340235, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.21852807999999999, + "after_init_synchronized_e2e_ms_per_call": 0.30511912, + "including_init_host_enqueue_ms_per_call": 0.56080715, + "including_init_synchronized_e2e_ms_per_call": 0.6485735500000002, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.148889608, + "after_init_synchronized_e2e_ms_per_call": 0.23568311200000003, + "including_init_host_enqueue_ms_per_call": 0.18311751500000004, + "including_init_synchronized_e2e_ms_per_call": 0.27002855500000006, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.87876, + "synchronized_e2e_ms": 7.94308, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.362336, + "median": 0.131264, + "min": 0.122368, + "p90": 0.1384, + "sample_count": 1061 + }, + "host_enqueue_ms": { + "max": 0.378817, + "median": 0.141152, + "min": 0.121952, + "p90": 0.168608, + "sample_count": 1061 + }, + "sample_count": 1061, + "synchronized_e2e_ms": { + "max": 0.458945, + "median": 0.227968, + "min": 0.209825, + "p90": 0.253057, + "sample_count": 1061 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 1290, + "candidate_precomputed_gpu_span_ms": 0.07744, + "candidate_precomputed_host_enqueue_ms": 0.037008, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.07744, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.037008, + "candidate_precomputed_synchronized_e2e_ms": 0.111265, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.07744 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.029696, + "synchronized_e2e_ms": 0.105536 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.07744 + }, + "host_enqueue_ms": { + "median": 0.037008 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.07744 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1290, + "submission_ms": { + "median": 0.037008 + }, + "synchronized_e2e_ms": { + "median": 0.111265 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01dba090", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01dbb560" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.218595041322314, + "submission": 1.0998973194984867, + "synchronized_e2e": 1.2283467397654249 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 1061, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.094368, + "candidate_public_raw_host_enqueue_ms": 0.040705, + "candidate_public_raw_inter_kernel_gap_ms": 0.000192, + "candidate_public_raw_kernel_sum_ms": 0.094176, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.040705, + "candidate_public_raw_synchronized_e2e_ms": 0.136672, + "candidate_public_raw_tflops_from_gpu_span": 44.80182299084436, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.094176 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.050816, + "synchronized_e2e_ms": 0.138752 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.094368 + }, + "host_enqueue_ms": { + "median": 0.040705 + }, + "inter_kernel_gap_ms": { + "median": 0.000192 + }, + "kernel_sum_ms": { + "median": 0.094176 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1061, + "submission_ms": { + "median": 0.040705 + }, + "synchronized_e2e_ms": { + "median": 0.136672 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 1.974882, + "after_init_synchronized_e2e_ms_per_call": 2.05613, + "including_init_host_enqueue_ms_per_call": 36.552422, + "including_init_synchronized_e2e_ms_per_call": 450.143922, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2341227, + "after_init_synchronized_e2e_ms_per_call": 0.32861779999999996, + "including_init_host_enqueue_ms_per_call": 3.6918767000000003, + "including_init_synchronized_e2e_ms_per_call": 45.137397, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06004677, + "after_init_synchronized_e2e_ms_per_call": 0.15586657999999998, + "including_init_host_enqueue_ms_per_call": 0.40582217, + "including_init_synchronized_e2e_ms_per_call": 4.6367445, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.04263917699999999, + "after_init_synchronized_e2e_ms_per_call": 0.138591458, + "including_init_host_enqueue_ms_per_call": 0.07721671699999999, + "including_init_synchronized_e2e_ms_per_call": 0.5866792499999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 1.974882, + "synchronized_e2e_ms": 2.05613, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.096, + "median": 0.094368, + "min": 0.089888, + "p90": 0.09504, + "sample_count": 1061 + }, + "host_enqueue_ms": { + "max": 0.529921, + "median": 0.040705, + "min": 0.033088, + "p90": 0.049216, + "sample_count": 1061 + }, + "sample_count": 1061, + "synchronized_e2e_ms": { + "max": 0.611233, + "median": 0.136672, + "min": 0.128416, + "p90": 0.14288, + "sample_count": 1061 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.060864, + "submission_ms": 0.060864, + "synchronized_e2e_ms": 0.141952 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.05872, + "submission_ms": 0.05872, + "synchronized_e2e_ms": 0.127648 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.184576, + "submission_ms": 0.184576, + "synchronized_e2e_ms": 0.264832 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.87876, + "submission_ms": 7.87876, + "synchronized_e2e_ms": 7.94308 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.050464, + "submission_ms": 0.050464, + "synchronized_e2e_ms": 0.119329 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.029696, + "submission_ms": 0.029696, + "synchronized_e2e_ms": 0.105536 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 0.976226, + "submission_ms": 0.976226, + "synchronized_e2e_ms": 1.000066 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 0.860576, + "submission_ms": 0.860576, + "synchronized_e2e_ms": 0.880192 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.050816, + "submission_ms": 0.050816, + "synchronized_e2e_ms": 0.138752 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 1.974882, + "submission_ms": 1.974882, + "synchronized_e2e_ms": 2.05613 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 1.518318, + "evolution_kernel_ms": 0.287872, + "evolution_speedup": 5.2743, + "evolution_tflops": 14.6866, + "expected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_d9d5_d112_highk_request_b3_n768_k8192_d112", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 4702, + "measurement_schedule_sha256": "b74eb4fe3ed0350f13df784f6e89efd328b79c233bb9406401634261bc4a984b", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 1290, + "public_pair_count": 1061, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 1290, + "baseline_public_raw": 1061, + "candidate_precomputed": 1290, + "candidate_public_raw": 1061 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 942 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:adjacent_d9d5_d112_highk_request_b3_n768_k8192_d112", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.3801652892561984, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.6679934441582769, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.863121495236196, + "including_init_synchronized_e2e_speedup": 0.09394444961538324, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 3.041463974258242, + "including_init_synchronized_e2e_speedup": 0.09823392119842446, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 1.9575660157552701, + "including_init_synchronized_e2e_speedup": 0.13987692226733653, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.7005601600641218, + "including_init_synchronized_e2e_speedup": 0.4602660738384733, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.39097999321804, + "hot_synchronized_e2e_speedup": 1.6679934441582769, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 9511201, + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_d9d5_d112_highk_request_b3_n768_k8192_d112", + "source": "request_specific", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 128, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 4, + "num_warps": 4 + } + }, + { + "B": 5, + "D": 128, + "K": 1024, + "N": 6016, + "baseline_07cf_adapter_bench_iters": 2454, + "baseline_07cf_adapter_gpu_span_ms": 0.07408, + "baseline_07cf_adapter_host_enqueue_ms": 0.15088, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.023344, + "baseline_07cf_adapter_kernel_sum_ms": 0.05072, + "baseline_07cf_adapter_submission_ms": 0.15088, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.17824, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.05072 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.198816, + "synchronized_e2e_ms": 0.222144 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.07408 + }, + "host_enqueue_ms": { + "median": 0.15088 + }, + "inter_kernel_gap_ms": { + "median": 0.023344 + }, + "kernel_sum_ms": { + "median": 0.05072 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2454, + "submission_ms": { + "median": 0.15088 + }, + "synchronized_e2e_ms": { + "median": 0.17824 + } + }, + "baseline_07cf_precomputed_bench_iters": 5140, + "baseline_07cf_precomputed_gpu_span_ms": 0.029248, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042752, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.029248, + "baseline_07cf_precomputed_submission_ms": 0.042752, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.079168, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.029248 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.043136, + "synchronized_e2e_ms": 0.078368 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.029248 + }, + "host_enqueue_ms": { + "median": 0.042752 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.029248 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5140, + "submission_ms": { + "median": 0.042752 + }, + "synchronized_e2e_ms": { + "median": 0.079168 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.5328227571115973, + "submission": 3.5291916167664668, + "synchronized_e2e": 2.2514147130153597 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.796393, + "after_init_synchronized_e2e_ms_per_call": 8.822281, + "including_init_host_enqueue_ms_per_call": 43.245837, + "including_init_synchronized_e2e_ms_per_call": 492.863775, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 1.0154313, + "after_init_synchronized_e2e_ms_per_call": 1.0426441, + "including_init_host_enqueue_ms_per_call": 4.4603757, + "including_init_synchronized_e2e_ms_per_call": 49.4467935, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23733512999999998, + "after_init_synchronized_e2e_ms_per_call": 0.26468041, + "including_init_host_enqueue_ms_per_call": 0.58182957, + "including_init_synchronized_e2e_ms_per_call": 5.10509535, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15952551299999998, + "after_init_synchronized_e2e_ms_per_call": 0.18688404100000003, + "including_init_host_enqueue_ms_per_call": 0.19397495699999998, + "including_init_synchronized_e2e_ms_per_call": 0.670925535, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.796393, + "synchronized_e2e_ms": 8.822281, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 7.94052, + "median": 0.07408, + "min": 0.065856, + "p90": 0.09481280000000007, + "sample_count": 2454 + }, + "host_enqueue_ms": { + "max": 87.39033, + "median": 0.15088, + "min": 0.126816, + "p90": 0.20452550000000003, + "sample_count": 2454 + }, + "sample_count": 2454, + "synchronized_e2e_ms": { + "max": 87.62537, + "median": 0.17824, + "min": 0.154976, + "p90": 0.22760070000000002, + "sample_count": 2454 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 5140, + "candidate_precomputed_gpu_span_ms": 0.019744, + "candidate_precomputed_host_enqueue_ms": 0.041984, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.019744, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.041984, + "candidate_precomputed_synchronized_e2e_ms": 0.05632, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.019744 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.033472, + "synchronized_e2e_ms": 0.049088 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.019744 + }, + "host_enqueue_ms": { + "median": 0.041984 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.019744 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5140, + "submission_ms": { + "median": 0.041984 + }, + "synchronized_e2e_ms": { + "median": 0.05632 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7edfdf10", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7edfcaa0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.0696920583468392, + "submission": 1.1040396341463417, + "synchronized_e2e": 1.5840997869318183 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 2454, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.040864, + "candidate_public_raw_host_enqueue_ms": 0.046352000000000004, + "candidate_public_raw_inter_kernel_gap_ms": 0.00016, + "candidate_public_raw_kernel_sum_ms": 0.040704, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.046352000000000004, + "candidate_public_raw_synchronized_e2e_ms": 0.0892165, + "candidate_public_raw_tflops_from_gpu_span": 192.96425998433833, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.040704 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.239136, + "synchronized_e2e_ms": 0.272128 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.040864 + }, + "host_enqueue_ms": { + "median": 0.046352000000000004 + }, + "inter_kernel_gap_ms": { + "median": 0.00016 + }, + "kernel_sum_ms": { + "median": 0.040704 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2454, + "submission_ms": { + "median": 0.046352000000000004 + }, + "synchronized_e2e_ms": { + "median": 0.0892165 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.930371, + "after_init_synchronized_e2e_ms_per_call": 2.956964, + "including_init_host_enqueue_ms_per_call": 37.804455, + "including_init_synchronized_e2e_ms_per_call": 37.90004, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.33475390000000005, + "after_init_synchronized_e2e_ms_per_call": 0.37599125000000005, + "including_init_host_enqueue_ms_per_call": 3.8221622999999996, + "including_init_synchronized_e2e_ms_per_call": 3.8702988499999997, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07519219, + "after_init_synchronized_e2e_ms_per_call": 0.117893975, + "including_init_host_enqueue_ms_per_call": 0.42393302999999993, + "including_init_synchronized_e2e_ms_per_call": 0.467324735, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.049236019000000006, + "after_init_synchronized_e2e_ms_per_call": 0.09208424750000001, + "including_init_host_enqueue_ms_per_call": 0.084110103, + "including_init_synchronized_e2e_ms_per_call": 0.12702732349999998, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.930371, + "synchronized_e2e_ms": 2.956964, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.041728, + "median": 0.040864, + "min": 0.040256, + "p90": 0.041152, + "sample_count": 2454 + }, + "host_enqueue_ms": { + "max": 0.717217, + "median": 0.046352000000000004, + "min": 0.036, + "p90": 0.06596800000000004, + "sample_count": 2454 + }, + "sample_count": 2454, + "synchronized_e2e_ms": { + "max": 0.746017, + "median": 0.0892165, + "min": 0.080032, + "p90": 0.1072544, + "sample_count": 2454 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.066144, + "submission_ms": 0.066144, + "synchronized_e2e_ms": 0.10128 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.043136, + "submission_ms": 0.043136, + "synchronized_e2e_ms": 0.078368 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.198816, + "submission_ms": 0.198816, + "synchronized_e2e_ms": 0.222144 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.796393, + "submission_ms": 8.796393, + "synchronized_e2e_ms": 8.822281 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.056801, + "submission_ms": 0.056801, + "synchronized_e2e_ms": 0.071137 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.033472, + "submission_ms": 0.033472, + "synchronized_e2e_ms": 0.049088 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.280385, + "submission_ms": 1.280385, + "synchronized_e2e_ms": 1.301153 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.052577, + "submission_ms": 1.052577, + "synchronized_e2e_ms": 1.070753 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.239136, + "submission_ms": 0.239136, + "synchronized_e2e_ms": 0.272128 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.930371, + "submission_ms": 2.930371, + "synchronized_e2e_ms": 2.956964 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.460207, + "evolution_kernel_ms": 0.159104, + "evolution_speedup": 2.8925, + "evolution_tflops": 49.5606, + "expected_route": "aligned_weave_v10_fallback", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_d9d5_d128_guard_miss_fallback_b5_n6016_k1024_d128", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 15188, + "measurement_schedule_sha256": "d692f5847d6ad2f1ac4943eab9549b2c86ad54eae986bf45922fda9e542d293b", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 5140, + "public_pair_count": 2454, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 5140, + "baseline_public_raw": 2454, + "candidate_precomputed": 5140, + "candidate_public_raw": 2454 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 3038 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:adjacent_d9d5_d128_guard_miss_fallback_b5_n6016_k1024_d128", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.4813614262560777, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.997836723027691, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.9835605032729515, + "including_init_synchronized_e2e_speedup": 13.00430751524273, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.7730541601699503, + "including_init_synchronized_e2e_speedup": 12.775962636580378, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.245071556879815, + "including_init_synchronized_e2e_speedup": 10.924085475595467, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.029489799544705, + "including_init_synchronized_e2e_speedup": 5.2817418844537025, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.8128425998433828, + "hot_synchronized_e2e_speedup": 1.997836723027691, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 9512801, + "selected_route": "aligned_weave_v10_fallback", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_d9d5_d128_guard_miss_fallback_b5_n6016_k1024_d128", + "source": "guard_miss_fallback", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 224, + "K": 1280, + "N": 2944, + "baseline_07cf_adapter_bench_iters": 3341, + "baseline_07cf_adapter_gpu_span_ms": 0.079136, + "baseline_07cf_adapter_host_enqueue_ms": 0.161248, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.035712, + "baseline_07cf_adapter_kernel_sum_ms": 0.043456, + "baseline_07cf_adapter_submission_ms": 0.161248, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.187104, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.043456 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.267616, + "synchronized_e2e_ms": 0.29696 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.079136 + }, + "host_enqueue_ms": { + "median": 0.161248 + }, + "inter_kernel_gap_ms": { + "median": 0.035712 + }, + "kernel_sum_ms": { + "median": 0.043456 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3341, + "submission_ms": { + "median": 0.161248 + }, + "synchronized_e2e_ms": { + "median": 0.187104 + } + }, + "baseline_07cf_precomputed_bench_iters": 4002, + "baseline_07cf_precomputed_gpu_span_ms": 0.038432, + "baseline_07cf_precomputed_host_enqueue_ms": 0.047008, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.038432, + "baseline_07cf_precomputed_submission_ms": 0.047008, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.0904485, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.038432 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.05856, + "synchronized_e2e_ms": 0.090784 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.038432 + }, + "host_enqueue_ms": { + "median": 0.047008 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.038432 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4002, + "submission_ms": { + "median": 0.047008 + }, + "synchronized_e2e_ms": { + "median": 0.0904485 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.0591174021648624, + "submission": 3.4302246426140233, + "synchronized_e2e": 2.0686246869765665 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.595559, + "after_init_synchronized_e2e_ms_per_call": 7.628136, + "including_init_host_enqueue_ms_per_call": 43.146827, + "including_init_synchronized_e2e_ms_per_call": 43.259373, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9046790999999998, + "after_init_synchronized_e2e_ms_per_call": 0.9312071999999999, + "including_init_host_enqueue_ms_per_call": 4.4598059, + "including_init_synchronized_e2e_ms_per_call": 4.4943309, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23559111000000002, + "after_init_synchronized_e2e_ms_per_call": 0.26151432, + "including_init_host_enqueue_ms_per_call": 0.59110379, + "including_init_synchronized_e2e_ms_per_call": 0.61782669, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.168682311, + "after_init_synchronized_e2e_ms_per_call": 0.19454503199999998, + "including_init_host_enqueue_ms_per_call": 0.204233579, + "including_init_synchronized_e2e_ms_per_call": 0.230176269, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.595559, + "synchronized_e2e_ms": 7.628136, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.397921, + "median": 0.079136, + "min": 0.068768, + "p90": 0.104833, + "sample_count": 3341 + }, + "host_enqueue_ms": { + "max": 22.803575, + "median": 0.161248, + "min": 0.126336, + "p90": 0.260192, + "sample_count": 3341 + }, + "sample_count": 3341, + "synchronized_e2e_ms": { + "max": 22.863575, + "median": 0.187104, + "min": 0.153824, + "p90": 0.289664, + "sample_count": 3341 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 4002, + "candidate_precomputed_gpu_span_ms": 0.022016, + "candidate_precomputed_host_enqueue_ms": 0.057024, + "candidate_precomputed_inter_kernel_gap_ms": 0.004064, + "candidate_precomputed_kernel_sum_ms": 0.017952, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.057024, + "candidate_precomputed_synchronized_e2e_ms": 0.069761, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.017952 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.056512, + "synchronized_e2e_ms": 0.077376 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.022016 + }, + "host_enqueue_ms": { + "median": 0.057024 + }, + "inter_kernel_gap_ms": { + "median": 0.004064 + }, + "kernel_sum_ms": { + "median": 0.017952 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4002, + "submission_ms": { + "median": 0.057024 + }, + "synchronized_e2e_ms": { + "median": 0.069761 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc029a1c40", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc0121f7d0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.3619186046511627, + "submission": 0.8507295173961841, + "synchronized_e2e": 1.126589355083786 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 3341, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.029984, + "candidate_public_raw_host_enqueue_ms": 0.048512, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.02992, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.048512, + "candidate_public_raw_synchronized_e2e_ms": 0.078592, + "candidate_public_raw_tflops_from_gpu_span": 112.60721451440769, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.02992 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.083072, + "synchronized_e2e_ms": 0.10896 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.029984 + }, + "host_enqueue_ms": { + "median": 0.048512 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.02992 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3341, + "submission_ms": { + "median": 0.048512 + }, + "synchronized_e2e_ms": { + "median": 0.078592 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.671363, + "after_init_synchronized_e2e_ms_per_call": 2.722691, + "including_init_host_enqueue_ms_per_call": 38.509383, + "including_init_synchronized_e2e_ms_per_call": 486.822904, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.3107971, + "after_init_synchronized_e2e_ms_per_call": 0.3430019, + "including_init_host_enqueue_ms_per_call": 3.8945990999999998, + "including_init_synchronized_e2e_ms_per_call": 48.7530232, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07474051, + "after_init_synchronized_e2e_ms_per_call": 0.10503298999999998, + "including_init_host_enqueue_ms_per_call": 0.43312071, + "including_init_synchronized_e2e_ms_per_call": 4.9460351199999995, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.051134850999999995, + "after_init_synchronized_e2e_ms_per_call": 0.08123609899999999, + "including_init_host_enqueue_ms_per_call": 0.086972871, + "including_init_synchronized_e2e_ms_per_call": 0.565336312, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.671363, + "synchronized_e2e_ms": 2.722691, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.030784, + "median": 0.029984, + "min": 0.029536, + "p90": 0.03024, + "sample_count": 3341 + }, + "host_enqueue_ms": { + "max": 0.468641, + "median": 0.048512, + "min": 0.037184, + "p90": 0.080705, + "sample_count": 3341 + }, + "sample_count": 3341, + "synchronized_e2e_ms": { + "max": 5.385126, + "median": 0.078592, + "min": 0.068192, + "p90": 0.108833, + "sample_count": 3341 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.085344, + "submission_ms": 0.085344, + "synchronized_e2e_ms": 0.115456 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.05856, + "submission_ms": 0.05856, + "synchronized_e2e_ms": 0.090784 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.267616, + "submission_ms": 0.267616, + "synchronized_e2e_ms": 0.29696 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.595559, + "submission_ms": 7.595559, + "synchronized_e2e_ms": 7.628136 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.079488, + "submission_ms": 0.079488, + "synchronized_e2e_ms": 0.103552 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.056512, + "submission_ms": 0.056512, + "synchronized_e2e_ms": 0.077376 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.50733, + "submission_ms": 1.50733, + "synchronized_e2e_ms": 1.53805 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.366849, + "submission_ms": 1.366849, + "synchronized_e2e_ms": 1.413281 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.083072, + "submission_ms": 0.083072, + "synchronized_e2e_ms": 0.10896 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.671363, + "submission_ms": 2.671363, + "synchronized_e2e_ms": 2.722691 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.488255, + "evolution_kernel_ms": 0.180384, + "evolution_speedup": 2.7068, + "evolution_tflops": 18.7179, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_d9d5_d224_random_midk_b2_n2944_k1280_d224", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 14686, + "measurement_schedule_sha256": "a0fc0398493502a4bfcd709416e1f0a221226657afd61d0b1024ca7c41ea26c8", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 4002, + "public_pair_count": 3341, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 4002, + "baseline_public_raw": 3341, + "candidate_precomputed": 4002, + "candidate_public_raw": 3341 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2940 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:adjacent_d9d5_d224_random_midk_b2_n2944_k1280_d224", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.7456395348837208, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.380700325732899, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.801689945719143, + "including_init_synchronized_e2e_speedup": 0.08886059518678685, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.7148747572535306, + "including_init_synchronized_e2e_speedup": 0.09218568624068424, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.4898302904639777, + "including_init_synchronized_e2e_speedup": 0.12491352669570209, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.3948101200674343, + "including_init_synchronized_e2e_speedup": 0.4071492740059478, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.639274279615795, + "hot_synchronized_e2e_speedup": 2.380700325732899, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 9522401, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_d9d5_d224_random_midk_b2_n2944_k1280_d224", + "source": "random_legal", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 288, + "K": 4096, + "N": 896, + "baseline_07cf_adapter_bench_iters": 3622, + "baseline_07cf_adapter_gpu_span_ms": 0.14832050000000002, + "baseline_07cf_adapter_host_enqueue_ms": 0.161424, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.037824, + "baseline_07cf_adapter_kernel_sum_ms": 0.110433, + "baseline_07cf_adapter_submission_ms": 0.161424, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.257888, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.110433 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.183616, + "synchronized_e2e_ms": 0.274688 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.14832050000000002 + }, + "host_enqueue_ms": { + "median": 0.161424 + }, + "inter_kernel_gap_ms": { + "median": 0.037824 + }, + "kernel_sum_ms": { + "median": 0.110433 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3622, + "submission_ms": { + "median": 0.161424 + }, + "synchronized_e2e_ms": { + "median": 0.257888 + } + }, + "baseline_07cf_precomputed_bench_iters": 4661, + "baseline_07cf_precomputed_gpu_span_ms": 0.115744, + "baseline_07cf_precomputed_host_enqueue_ms": 0.046272, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.115744, + "baseline_07cf_precomputed_submission_ms": 0.046272, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.167456, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.115744 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.04704, + "synchronized_e2e_ms": 0.146272 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.115744 + }, + "host_enqueue_ms": { + "median": 0.046272 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.115744 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4661, + "submission_ms": { + "median": 0.046272 + }, + "synchronized_e2e_ms": { + "median": 0.167456 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.2814530342825547, + "submission": 3.4885892116182573, + "synchronized_e2e": 1.5400343970953565 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.785448, + "after_init_synchronized_e2e_ms_per_call": 7.86548, + "including_init_host_enqueue_ms_per_call": 44.719886, + "including_init_synchronized_e2e_ms_per_call": 462.976799, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9238263999999999, + "after_init_synchronized_e2e_ms_per_call": 1.0186472, + "including_init_host_enqueue_ms_per_call": 4.6172702, + "including_init_synchronized_e2e_ms_per_call": 46.5297791, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23766424, + "after_init_synchronized_e2e_ms_per_call": 0.33396391999999997, + "including_init_host_enqueue_ms_per_call": 0.60700862, + "including_init_synchronized_e2e_ms_per_call": 4.88507711, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.16904802400000002, + "after_init_synchronized_e2e_ms_per_call": 0.265495592, + "including_init_host_enqueue_ms_per_call": 0.20598246200000003, + "including_init_synchronized_e2e_ms_per_call": 0.7206069110000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.785448, + "synchronized_e2e_ms": 7.86548, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.33127, + "median": 0.14832050000000002, + "min": 0.137345, + "p90": 0.17496319999999999, + "sample_count": 3622 + }, + "host_enqueue_ms": { + "max": 47.892561, + "median": 0.161424, + "min": 0.12416, + "p90": 0.25543360000000004, + "sample_count": 3622 + }, + "sample_count": 3622, + "synchronized_e2e_ms": { + "max": 47.942929, + "median": 0.257888, + "min": 0.224992, + "p90": 0.3403817, + "sample_count": 3622 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 4661, + "candidate_precomputed_gpu_span_ms": 0.02208, + "candidate_precomputed_host_enqueue_ms": 0.053824, + "candidate_precomputed_inter_kernel_gap_ms": 0.002208, + "candidate_precomputed_kernel_sum_ms": 0.01984, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.053824, + "candidate_precomputed_synchronized_e2e_ms": 0.066752, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.01984 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.035936, + "synchronized_e2e_ms": 0.050112 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.02208 + }, + "host_enqueue_ms": { + "median": 0.053824 + }, + "inter_kernel_gap_ms": { + "median": 0.002208 + }, + "kernel_sum_ms": { + "median": 0.01984 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4661, + "submission_ms": { + "median": 0.053824 + }, + "synchronized_e2e_ms": { + "median": 0.066752 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc03ef40b0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc03ef5130" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.2608695652173914, + "submission": 0.9771105826397147, + "synchronized_e2e": 1.2176414189837008 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 3622, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.02784, + "candidate_public_raw_host_enqueue_ms": 0.052592, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.027872, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.052592, + "candidate_public_raw_synchronized_e2e_ms": 0.08128, + "candidate_public_raw_tflops_from_gpu_span": 75.93136551724137, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.027776 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.239008, + "synchronized_e2e_ms": 0.260832 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.02784 + }, + "host_enqueue_ms": { + "median": 0.052592 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.027872 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3622, + "submission_ms": { + "median": 0.052592 + }, + "synchronized_e2e_ms": { + "median": 0.08128 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.418531, + "after_init_synchronized_e2e_ms_per_call": 2.444291, + "including_init_host_enqueue_ms_per_call": 39.776554, + "including_init_synchronized_e2e_ms_per_call": 39.879914, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2891859, + "after_init_synchronized_e2e_ms_per_call": 0.31758110000000006, + "including_init_host_enqueue_ms_per_call": 4.0249882, + "including_init_synchronized_e2e_ms_per_call": 4.061143400000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07625139, + "after_init_synchronized_e2e_ms_per_call": 0.10491011, + "including_init_host_enqueue_ms_per_call": 0.44983162, + "including_init_synchronized_e2e_ms_per_call": 0.47926634, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.054957939000000004, + "after_init_synchronized_e2e_ms_per_call": 0.08364301100000002, + "including_init_host_enqueue_ms_per_call": 0.092315962, + "including_init_synchronized_e2e_ms_per_call": 0.121078634, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.418531, + "synchronized_e2e_ms": 2.444291, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.028736, + "median": 0.02784, + "min": 0.026688, + "p90": 0.02832, + "sample_count": 3622 + }, + "host_enqueue_ms": { + "max": 15.469168, + "median": 0.052592, + "min": 0.038177, + "p90": 0.08481600000000004, + "sample_count": 3622 + }, + "sample_count": 3622, + "synchronized_e2e_ms": { + "max": 20.614773, + "median": 0.08128, + "min": 0.068577, + "p90": 0.11015680000000001, + "sample_count": 3622 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.063232, + "submission_ms": 0.063232, + "synchronized_e2e_ms": 0.159873 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.04704, + "submission_ms": 0.04704, + "synchronized_e2e_ms": 0.146272 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.183616, + "submission_ms": 0.183616, + "synchronized_e2e_ms": 0.274688 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.785448, + "submission_ms": 7.785448, + "synchronized_e2e_ms": 7.86548 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.064544, + "submission_ms": 0.064544, + "synchronized_e2e_ms": 0.08064 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.035936, + "submission_ms": 0.035936, + "synchronized_e2e_ms": 0.050112 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.315042, + "submission_ms": 1.315042, + "synchronized_e2e_ms": 1.335426 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.241057, + "submission_ms": 1.241057, + "synchronized_e2e_ms": 1.263745 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.239008, + "submission_ms": 0.239008, + "synchronized_e2e_ms": 0.260832 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.418531, + "submission_ms": 2.418531, + "synchronized_e2e_ms": 2.444291 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 57.939113, + "evolution_kernel_ms": 0.2448, + "evolution_speedup": 236.6794, + "evolution_tflops": 8.6353, + "expected_route": "d288_parent_splitk_hybrid_20260629_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_d9d5_d288_heldout_low_n_b1_n896_k4096_d288", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 16566, + "measurement_schedule_sha256": "7127c4a1052a0c7e0de4fe010b3afa3eb69fcb0aaece79029b3016c388090845", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 4661, + "public_pair_count": 3622, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 4661, + "baseline_public_raw": 3622, + "candidate_precomputed": 4661, + "candidate_public_raw": 3622 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 3316 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:adjacent_d9d5_d288_heldout_low_n_b1_n896_k4096_d288", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 5.242028985507247, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 3.1728346456692913, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.2178983598925, + "including_init_synchronized_e2e_speedup": 11.609272753196008, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 3.2075183315379907, + "including_init_synchronized_e2e_speedup": 11.457310052139501, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 3.1833339989825573, + "including_init_synchronized_e2e_speedup": 10.1928232848566, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 3.1741515378971705, + "including_init_synchronized_e2e_speedup": 5.951561288674599, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 5.327604166666667, + "hot_synchronized_e2e_speedup": 3.1728346456692913, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 9528801, + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_d9d5_d288_heldout_low_n_b1_n896_k4096_d288", + "source": "heldout_neighborhood", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 5, + "D": 352, + "K": 768, + "N": 2304, + "baseline_07cf_adapter_bench_iters": 1994, + "baseline_07cf_adapter_gpu_span_ms": 0.088832, + "baseline_07cf_adapter_host_enqueue_ms": 0.1432, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.018608, + "baseline_07cf_adapter_kernel_sum_ms": 0.0704, + "baseline_07cf_adapter_submission_ms": 0.1432, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.1861285, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.0704 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.193056, + "synchronized_e2e_ms": 0.23024 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.088832 + }, + "host_enqueue_ms": { + "median": 0.1432 + }, + "inter_kernel_gap_ms": { + "median": 0.018608 + }, + "kernel_sum_ms": { + "median": 0.0704 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1994, + "submission_ms": { + "median": 0.1432 + }, + "synchronized_e2e_ms": { + "median": 0.1861285 + } + }, + "baseline_07cf_precomputed_bench_iters": 3217, + "baseline_07cf_precomputed_gpu_span_ms": 0.05312, + "baseline_07cf_precomputed_host_enqueue_ms": 0.041568, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.05312, + "baseline_07cf_precomputed_submission_ms": 0.041568, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.101536, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.05312 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.04528, + "synchronized_e2e_ms": 0.093664 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.05312 + }, + "host_enqueue_ms": { + "median": 0.041568 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.05312 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3217, + "submission_ms": { + "median": 0.041568 + }, + "synchronized_e2e_ms": { + "median": 0.101536 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.672289156626506, + "submission": 3.44495765973826, + "synchronized_e2e": 1.8331281515915538 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 6.981768, + "after_init_synchronized_e2e_ms_per_call": 7.008584, + "including_init_host_enqueue_ms_per_call": 41.209675000000004, + "including_init_synchronized_e2e_ms_per_call": 41.354027, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8270567999999999, + "after_init_synchronized_e2e_ms_per_call": 0.8683740499999999, + "including_init_host_enqueue_ms_per_call": 4.2498475000000004, + "including_init_synchronized_e2e_ms_per_call": 4.302918350000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.21158568, + "after_init_synchronized_e2e_ms_per_call": 0.254353055, + "including_init_host_enqueue_ms_per_call": 0.55386475, + "including_init_synchronized_e2e_ms_per_call": 0.597807485, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15003856799999998, + "after_init_synchronized_e2e_ms_per_call": 0.19295095550000002, + "including_init_host_enqueue_ms_per_call": 0.18426647499999999, + "including_init_synchronized_e2e_ms_per_call": 0.2272963985, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 6.981768, + "synchronized_e2e_ms": 7.008584, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.368769, + "median": 0.088832, + "min": 0.079712, + "p90": 0.10745950000000001, + "sample_count": 1994 + }, + "host_enqueue_ms": { + "max": 46.816848, + "median": 0.1432, + "min": 0.123424, + "p90": 0.17517760000000002, + "sample_count": 1994 + }, + "sample_count": 1994, + "synchronized_e2e_ms": { + "max": 46.958641, + "median": 0.1861285, + "min": 0.166592, + "p90": 0.2155847, + "sample_count": 1994 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3217, + "candidate_precomputed_gpu_span_ms": 0.028288, + "candidate_precomputed_host_enqueue_ms": 0.052384, + "candidate_precomputed_inter_kernel_gap_ms": 0.001824, + "candidate_precomputed_kernel_sum_ms": 0.026464, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.052384, + "candidate_precomputed_synchronized_e2e_ms": 0.064416, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.026464 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.04592, + "synchronized_e2e_ms": 0.062208 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.028288 + }, + "host_enqueue_ms": { + "median": 0.052384 + }, + "inter_kernel_gap_ms": { + "median": 0.001824 + }, + "kernel_sum_ms": { + "median": 0.026464 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3217, + "submission_ms": { + "median": 0.052384 + }, + "synchronized_e2e_ms": { + "median": 0.064416 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb9355ddf0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb9355c4d0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.7726244343891402, + "submission": 0.8344532681734881, + "synchronized_e2e": 1.4848484848484849 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 1994, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.050144, + "candidate_public_raw_host_enqueue_ms": 0.043712, + "candidate_public_raw_inter_kernel_gap_ms": 0.000128, + "candidate_public_raw_kernel_sum_ms": 0.050016, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.043712, + "candidate_public_raw_synchronized_e2e_ms": 0.095648, + "candidate_public_raw_tflops_from_gpu_span": 124.21309508615188, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.050016 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.072448, + "synchronized_e2e_ms": 0.1048 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.050144 + }, + "host_enqueue_ms": { + "median": 0.043712 + }, + "inter_kernel_gap_ms": { + "median": 0.000128 + }, + "kernel_sum_ms": { + "median": 0.050016 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1994, + "submission_ms": { + "median": 0.043712 + }, + "synchronized_e2e_ms": { + "median": 0.095648 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.464546, + "after_init_synchronized_e2e_ms_per_call": 2.49165, + "including_init_host_enqueue_ms_per_call": 37.042086, + "including_init_synchronized_e2e_ms_per_call": 450.579442, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2857954, + "after_init_synchronized_e2e_ms_per_call": 0.3352482, + "including_init_host_enqueue_ms_per_call": 3.7435494, + "including_init_synchronized_e2e_ms_per_call": 45.1440274, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06792034, + "after_init_synchronized_e2e_ms_per_call": 0.11960802, + "including_init_host_enqueue_ms_per_call": 0.41369574, + "including_init_synchronized_e2e_ms_per_call": 4.6004859399999996, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.046132834000000005, + "after_init_synchronized_e2e_ms_per_call": 0.098044002, + "including_init_host_enqueue_ms_per_call": 0.080710374, + "including_init_synchronized_e2e_ms_per_call": 0.546131794, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.464546, + "synchronized_e2e_ms": 2.49165, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.050976, + "median": 0.050144, + "min": 0.04976, + "p90": 0.050465, + "sample_count": 1994 + }, + "host_enqueue_ms": { + "max": 0.581889, + "median": 0.043712, + "min": 0.033504, + "p90": 0.062083200000000005, + "sample_count": 1994 + }, + "sample_count": 1994, + "synchronized_e2e_ms": { + "max": 0.682721, + "median": 0.095648, + "min": 0.087809, + "p90": 0.11267200000000001, + "sample_count": 1994 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.065632, + "submission_ms": 0.065632, + "synchronized_e2e_ms": 0.112 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.04528, + "submission_ms": 0.04528, + "synchronized_e2e_ms": 0.093664 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.193056, + "submission_ms": 0.193056, + "synchronized_e2e_ms": 0.23024 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 6.981768, + "submission_ms": 6.981768, + "synchronized_e2e_ms": 7.008584 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.085888, + "submission_ms": 0.085888, + "synchronized_e2e_ms": 0.103584 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.04592, + "submission_ms": 0.04592, + "synchronized_e2e_ms": 0.062208 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.253217, + "submission_ms": 1.253217, + "synchronized_e2e_ms": 1.275233 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.039809, + "submission_ms": 1.039809, + "synchronized_e2e_ms": 1.060097 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.072448, + "submission_ms": 0.072448, + "synchronized_e2e_ms": 0.1048 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.464546, + "submission_ms": 2.464546, + "synchronized_e2e_ms": 2.49165 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.582592, + "evolution_kernel_ms": 0.186944, + "evolution_speedup": 3.1164, + "evolution_tflops": 33.3177, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_d9d5_d352_random_b5_n2304_k768_d352", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 10422, + "measurement_schedule_sha256": "a357c7403b036e05c1ff148eefd346500740abfa585e79534ec1afdb75e9f773", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3217, + "public_pair_count": 1994, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3217, + "baseline_public_raw": 1994, + "candidate_precomputed": 3217, + "candidate_public_raw": 1994 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2086 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:adjacent_d9d5_d352_random_b5_n2304_k768_d352", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.8778280542986425, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.9459737788558047, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.8128284470130236, + "including_init_synchronized_e2e_speedup": 0.09177965780338465, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.5902422444028033, + "including_init_synchronized_e2e_speedup": 0.09531534065124196, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.126555184175777, + "including_init_synchronized_e2e_speedup": 0.12994442169733053, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.9680036673737575, + "including_init_synchronized_e2e_speedup": 0.4161933090092169, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.7715379706445435, + "hot_synchronized_e2e_speedup": 1.9459737788558047, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 9535201, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_d9d5_d352_random_b5_n2304_k768_d352", + "source": "random_legal", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 416, + "K": 8192, + "N": 640, + "baseline_07cf_adapter_bench_iters": 1424, + "baseline_07cf_adapter_gpu_span_ms": 0.2588165, + "baseline_07cf_adapter_host_enqueue_ms": 0.15184, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.016992, + "baseline_07cf_adapter_kernel_sum_ms": 0.241824, + "baseline_07cf_adapter_submission_ms": 0.15184, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.36420850000000005, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.241824 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.28048, + "synchronized_e2e_ms": 0.48544 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.2588165 + }, + "host_enqueue_ms": { + "median": 0.15184 + }, + "inter_kernel_gap_ms": { + "median": 0.016992 + }, + "kernel_sum_ms": { + "median": 0.241824 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1424, + "submission_ms": { + "median": 0.15184 + }, + "synchronized_e2e_ms": { + "median": 0.36420850000000005 + } + }, + "baseline_07cf_precomputed_bench_iters": 1899, + "baseline_07cf_precomputed_gpu_span_ms": 0.242336, + "baseline_07cf_precomputed_host_enqueue_ms": 0.043744, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.242336, + "baseline_07cf_precomputed_submission_ms": 0.043744, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.293729, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.242336 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.04784, + "synchronized_e2e_ms": 0.265889 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.242336 + }, + "host_enqueue_ms": { + "median": 0.043744 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.242336 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1899, + "submission_ms": { + "median": 0.043744 + }, + "synchronized_e2e_ms": { + "median": 0.293729 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.0680068169813812, + "submission": 3.471104608632041, + "synchronized_e2e": 1.2399473664500271 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.137193, + "after_init_synchronized_e2e_ms_per_call": 8.327561, + "including_init_host_enqueue_ms_per_call": 42.586636999999996, + "including_init_synchronized_e2e_ms_per_call": 492.369055, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9503752999999999, + "after_init_synchronized_e2e_ms_per_call": 1.16054375, + "including_init_host_enqueue_ms_per_call": 4.3953197, + "including_init_synchronized_e2e_ms_per_call": 49.56469315, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23169353, + "after_init_synchronized_e2e_ms_per_call": 0.44384202500000003, + "including_init_host_enqueue_ms_per_call": 0.57618797, + "including_init_synchronized_e2e_ms_per_call": 5.284256965, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.159825353, + "after_init_synchronized_e2e_ms_per_call": 0.37217185250000007, + "including_init_host_enqueue_ms_per_call": 0.194274797, + "including_init_synchronized_e2e_ms_per_call": 0.8562133465, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.137193, + "synchronized_e2e_ms": 8.327561, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.752769, + "median": 0.2588165, + "min": 0.251616, + "p90": 0.2717408, + "sample_count": 1424 + }, + "host_enqueue_ms": { + "max": 0.795201, + "median": 0.15184, + "min": 0.131009, + "p90": 0.1799904, + "sample_count": 1424 + }, + "sample_count": 1424, + "synchronized_e2e_ms": { + "max": 1.073953, + "median": 0.36420850000000005, + "min": 0.343425, + "p90": 0.38965120000000003, + "sample_count": 1424 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 4, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 1899, + "candidate_precomputed_gpu_span_ms": 0.051392, + "candidate_precomputed_host_enqueue_ms": 0.06992, + "candidate_precomputed_inter_kernel_gap_ms": 0.004736, + "candidate_precomputed_kernel_sum_ms": 0.046432, + "candidate_precomputed_launch_count": 3, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.06992, + "candidate_precomputed_synchronized_e2e_ms": 0.089376, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.046432 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.059392, + "synchronized_e2e_ms": 0.081664 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 3.0 + }, + "gpu_span_ms": { + "median": 0.051392 + }, + "host_enqueue_ms": { + "median": 0.06992 + }, + "inter_kernel_gap_ms": { + "median": 0.004736 + }, + "kernel_sum_ms": { + "median": 0.046432 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1899, + "submission_ms": { + "median": 0.06992 + }, + "synchronized_e2e_ms": { + "median": 0.089376 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc02eb8ce0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc02eb8e00" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.3648816936488168, + "submission": 0.7068649885583523, + "synchronized_e2e": 1.3730755460078767 + }, + "candidate_public_raw_assignment_launch_count": 3, + "candidate_public_raw_bench_iters": 1424, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.070144, + "candidate_public_raw_host_enqueue_ms": 0.049423999999999996, + "candidate_public_raw_inter_kernel_gap_ms": 3.2e-05, + "candidate_public_raw_kernel_sum_ms": 0.070336, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.049423999999999996, + "candidate_public_raw_synchronized_e2e_ms": 0.12272, + "candidate_public_raw_tflops_from_gpu_span": 124.37489051094892, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.070112 + }, + "activity_count": { + "median": 4.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.089664, + "synchronized_e2e_ms": 0.142688 + }, + "correlated_kernel_activity_count": { + "median": 4.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.070144 + }, + "host_enqueue_ms": { + "median": 0.049423999999999996 + }, + "inter_kernel_gap_ms": { + "median": 3.2e-05 + }, + "kernel_sum_ms": { + "median": 0.070336 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1424, + "submission_ms": { + "median": 0.049423999999999996 + }, + "synchronized_e2e_ms": { + "median": 0.12272 + } + }, + "candidate_public_raw_total_launch_count": 4, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.813956, + "after_init_synchronized_e2e_ms_per_call": 3.8481, + "including_init_host_enqueue_ms_per_call": 38.688039999999994, + "including_init_synchronized_e2e_ms_per_call": 38.791176, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.42587720000000007, + "after_init_synchronized_e2e_ms_per_call": 0.49525800000000003, + "including_init_host_enqueue_ms_per_call": 3.9132855999999996, + "including_init_synchronized_e2e_ms_per_call": 3.9895656, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.08706932, + "after_init_synchronized_e2e_ms_per_call": 0.1599738, + "including_init_host_enqueue_ms_per_call": 0.4358101599999999, + "including_init_synchronized_e2e_ms_per_call": 0.50940456, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.053188532, + "after_init_synchronized_e2e_ms_per_call": 0.12644538, + "including_init_host_enqueue_ms_per_call": 0.088062616, + "including_init_synchronized_e2e_ms_per_call": 0.16138845599999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.813956, + "synchronized_e2e_ms": 3.8481, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.072352, + "median": 0.070144, + "min": 0.06896, + "p90": 0.071136, + "sample_count": 1424 + }, + "host_enqueue_ms": { + "max": 0.438753, + "median": 0.049423999999999996, + "min": 0.03904, + "p90": 0.06254080000000001, + "sample_count": 1424 + }, + "sample_count": 1424, + "synchronized_e2e_ms": { + "max": 0.527936, + "median": 0.12272, + "min": 0.112448, + "p90": 0.1344864, + "sample_count": 1424 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.074848, + "submission_ms": 0.074848, + "synchronized_e2e_ms": 0.298496 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.04784, + "submission_ms": 0.04784, + "synchronized_e2e_ms": 0.265889 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.28048, + "submission_ms": 0.28048, + "synchronized_e2e_ms": 0.48544 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.137193, + "submission_ms": 8.137193, + "synchronized_e2e_ms": 8.327561 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.098112, + "submission_ms": 0.098112, + "synchronized_e2e_ms": 0.115488 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.059392, + "submission_ms": 0.059392, + "synchronized_e2e_ms": 0.081664 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.753826, + "submission_ms": 1.753826, + "synchronized_e2e_ms": 1.777506 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.456034, + "submission_ms": 1.456034, + "synchronized_e2e_ms": 1.479682 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.089664, + "submission_ms": 0.089664, + "synchronized_e2e_ms": 0.142688 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.813956, + "submission_ms": 3.813956, + "synchronized_e2e_ms": 3.8481 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 3.254428, + "evolution_kernel_ms": 0.298432, + "evolution_speedup": 10.9051, + "evolution_tflops": 29.2333, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_d9d5_d416_highk_request_b2_n640_k8192_d416", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 6646, + "measurement_schedule_sha256": "7d4e2ad12cfb6f7530a159586e8f630cb27bdf283a873fb53723714235fe1413", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 1899, + "public_pair_count": 1424, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 1899, + "baseline_public_raw": 1424, + "candidate_precomputed": 1899, + "candidate_public_raw": 1424 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 1330 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:adjacent_d9d5_d416_highk_request_b2_n640_k8192_d416", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 4.7154420921544205, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.967800684485007, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.1640708401548814, + "including_init_synchronized_e2e_speedup": 12.69281073097655, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.343311465943003, + "including_init_synchronized_e2e_speedup": 12.423581442049729, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.774466975217192, + "including_init_synchronized_e2e_speedup": 10.373399415584345, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.9433408519947513, + "including_init_synchronized_e2e_speedup": 5.305294862601573, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 3.689788150091241, + "hot_synchronized_e2e_speedup": 2.967800684485007, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 9541601, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_d9d5_d416_highk_request_b2_n640_k8192_d416", + "source": "request_specific", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 3, + "D": 480, + "K": 4096, + "N": 1024, + "baseline_07cf_adapter_bench_iters": 967, + "baseline_07cf_adapter_gpu_span_ms": 0.160864, + "baseline_07cf_adapter_host_enqueue_ms": 0.161536, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.023392, + "baseline_07cf_adapter_kernel_sum_ms": 0.137505, + "baseline_07cf_adapter_submission_ms": 0.161536, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.270624, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.137505 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.30624, + "synchronized_e2e_ms": 0.406272 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.160864 + }, + "host_enqueue_ms": { + "median": 0.161536 + }, + "inter_kernel_gap_ms": { + "median": 0.023392 + }, + "kernel_sum_ms": { + "median": 0.137505 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 967, + "submission_ms": { + "median": 0.161536 + }, + "synchronized_e2e_ms": { + "median": 0.270624 + } + }, + "baseline_07cf_precomputed_bench_iters": 1218, + "baseline_07cf_precomputed_gpu_span_ms": 0.1336, + "baseline_07cf_precomputed_host_enqueue_ms": 0.048192, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.1336, + "baseline_07cf_precomputed_submission_ms": 0.048192, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.187456, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.1336 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.079072, + "synchronized_e2e_ms": 0.18288 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.1336 + }, + "host_enqueue_ms": { + "median": 0.048192 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.1336 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1218, + "submission_ms": { + "median": 0.048192 + }, + "synchronized_e2e_ms": { + "median": 0.187456 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.2040718562874253, + "submission": 3.3519256308100935, + "synchronized_e2e": 1.4436667804711503 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.886793, + "after_init_synchronized_e2e_ms_per_call": 8.960809, + "including_init_host_enqueue_ms_per_call": 44.438061000000005, + "including_init_synchronized_e2e_ms_per_call": 44.592045999999996, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 1.0340617, + "after_init_synchronized_e2e_ms_per_call": 1.1396425, + "including_init_host_enqueue_ms_per_call": 4.589188500000001, + "including_init_synchronized_e2e_ms_per_call": 4.702766199999999, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.24878857000000004, + "after_init_synchronized_e2e_ms_per_call": 0.35752585, + "including_init_host_enqueue_ms_per_call": 0.60430125, + "including_init_synchronized_e2e_ms_per_call": 0.71383822, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.17026125700000003, + "after_init_synchronized_e2e_ms_per_call": 0.27931418499999994, + "including_init_host_enqueue_ms_per_call": 0.20581252500000002, + "including_init_synchronized_e2e_ms_per_call": 0.314945422, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.886793, + "synchronized_e2e_ms": 8.960809, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.561697, + "median": 0.160864, + "min": 0.150144, + "p90": 0.16906300000000002, + "sample_count": 967 + }, + "host_enqueue_ms": { + "max": 0.618784, + "median": 0.161536, + "min": 0.131648, + "p90": 0.18549759999999998, + "sample_count": 967 + }, + "sample_count": 967, + "synchronized_e2e_ms": { + "max": 0.704769, + "median": 0.270624, + "min": 0.241696, + "p90": 0.29221759999999997, + "sample_count": 967 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 1218, + "candidate_precomputed_gpu_span_ms": 0.082816, + "candidate_precomputed_host_enqueue_ms": 0.0554245, + "candidate_precomputed_inter_kernel_gap_ms": 0.002112, + "candidate_precomputed_kernel_sum_ms": 0.080672, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.0554245, + "candidate_precomputed_synchronized_e2e_ms": 0.1211525, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.080672 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.059776, + "synchronized_e2e_ms": 0.10384 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.082816 + }, + "host_enqueue_ms": { + "median": 0.0554245 + }, + "inter_kernel_gap_ms": { + "median": 0.002112 + }, + "kernel_sum_ms": { + "median": 0.080672 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1218, + "submission_ms": { + "median": 0.0554245 + }, + "synchronized_e2e_ms": { + "median": 0.1211525 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb937d1730", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb937d0620" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.2523183925811436, + "submission": 1.0069193226822073, + "synchronized_e2e": 1.3201213346814964 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 967, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.103712, + "candidate_public_raw_host_enqueue_ms": 0.055808, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.103648, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.055808, + "candidate_public_raw_synchronized_e2e_ms": 0.159936, + "candidate_public_raw_tflops_from_gpu_span": 116.4724961431657, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.103648 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.306304, + "synchronized_e2e_ms": 0.38288 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.103712 + }, + "host_enqueue_ms": { + "median": 0.055808 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.103648 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 967, + "submission_ms": { + "median": 0.055808 + }, + "synchronized_e2e_ms": { + "median": 0.159936 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 59.324542, + "after_init_synchronized_e2e_ms_per_call": 59.391102, + "including_init_host_enqueue_ms_per_call": 95.16256200000001, + "including_init_synchronized_e2e_ms_per_call": 543.491315, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 5.9826814, + "after_init_synchronized_e2e_ms_per_call": 6.0830526, + "including_init_host_enqueue_ms_per_call": 9.566483400000001, + "including_init_synchronized_e2e_ms_per_call": 54.4930739, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.6484953400000001, + "after_init_synchronized_e2e_ms_per_call": 0.7522476599999999, + "including_init_host_enqueue_ms_per_call": 1.00687554, + "including_init_synchronized_e2e_ms_per_call": 5.59324979, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.115076734, + "after_init_synchronized_e2e_ms_per_call": 0.21916716599999997, + "including_init_host_enqueue_ms_per_call": 0.150914754, + "including_init_synchronized_e2e_ms_per_call": 0.703267379, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 59.324542, + "synchronized_e2e_ms": 59.391102, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.106496, + "median": 0.103712, + "min": 0.100864, + "p90": 0.10544, + "sample_count": 967 + }, + "host_enqueue_ms": { + "max": 0.433152, + "median": 0.055808, + "min": 0.040001, + "p90": 0.0684736, + "sample_count": 967 + }, + "sample_count": 967, + "synchronized_e2e_ms": { + "max": 0.590497, + "median": 0.159936, + "min": 0.145473, + "p90": 0.17329279999999997, + "sample_count": 967 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.098432, + "submission_ms": 0.098432, + "synchronized_e2e_ms": 0.202144 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.079072, + "submission_ms": 0.079072, + "synchronized_e2e_ms": 0.18288 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.30624, + "submission_ms": 0.30624, + "synchronized_e2e_ms": 0.406272 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.886793, + "submission_ms": 8.886793, + "synchronized_e2e_ms": 8.960809 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.087008, + "submission_ms": 0.087008, + "synchronized_e2e_ms": 0.121792 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.059776, + "submission_ms": 0.059776, + "synchronized_e2e_ms": 0.10384 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.623266, + "submission_ms": 1.623266, + "synchronized_e2e_ms": 1.658114 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.467138, + "submission_ms": 1.467138, + "synchronized_e2e_ms": 1.502562 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.306304, + "submission_ms": 0.306304, + "synchronized_e2e_ms": 0.38288 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 59.324542, + "submission_ms": 59.324542, + "synchronized_e2e_ms": 59.391102 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 1.92259, + "evolution_kernel_ms": 0.297632, + "evolution_speedup": 6.4596, + "evolution_tflops": 40.5857, + "expected_route": "d480_splitk_k1024_eac2_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_d9d5_d480_heldout_highk_b3_n1024_k4096_d480", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 4370, + "measurement_schedule_sha256": "cbaba6718d85ce70509f2f85ae56ca526da9500fd1d67c02dc9bb1f6c8a2e85a", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 1218, + "public_pair_count": 967, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 1218, + "baseline_public_raw": 967, + "candidate_precomputed": 1218, + "candidate_public_raw": 967 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 876 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:adjacent_d9d5_d480_heldout_highk_b3_n1024_k4096_d480", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.6132148377125193, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.6920768307322929, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.15087797158570992, + "including_init_synchronized_e2e_speedup": 0.0820473938944176, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.18734713883618231, + "including_init_synchronized_e2e_speedup": 0.08630025548989996, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.47527678583938704, + "including_init_synchronized_e2e_speedup": 0.12762494914427913, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.2744344424292093, + "including_init_synchronized_e2e_speedup": 0.4478316944656692, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.5510644862696699, + "hot_synchronized_e2e_speedup": 1.6920768307322929, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 9548001, + "selected_route": "d480_splitk_k1024_eac2_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_d9d5_d480_heldout_highk_b3_n1024_k4096_d480", + "source": "heldout_neighborhood", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 48, + "K": 4096, + "N": 640, + "baseline_07cf_adapter_bench_iters": 2875, + "baseline_07cf_adapter_gpu_span_ms": 0.085824, + "baseline_07cf_adapter_host_enqueue_ms": 0.148992, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.0376, + "baseline_07cf_adapter_kernel_sum_ms": 0.048192, + "baseline_07cf_adapter_submission_ms": 0.148992, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.186721, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.048192 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.19568, + "synchronized_e2e_ms": 0.23056 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.085824 + }, + "host_enqueue_ms": { + "median": 0.148992 + }, + "inter_kernel_gap_ms": { + "median": 0.0376 + }, + "kernel_sum_ms": { + "median": 0.048192 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2875, + "submission_ms": { + "median": 0.148992 + }, + "synchronized_e2e_ms": { + "median": 0.186721 + } + }, + "baseline_07cf_precomputed_bench_iters": 2567, + "baseline_07cf_precomputed_gpu_span_ms": 0.04992, + "baseline_07cf_precomputed_host_enqueue_ms": 0.045024, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.04992, + "baseline_07cf_precomputed_submission_ms": 0.045024, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.100288, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.04992 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.045024, + "synchronized_e2e_ms": 0.089312 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.04992 + }, + "host_enqueue_ms": { + "median": 0.045024 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.04992 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2567, + "submission_ms": { + "median": 0.045024 + }, + "synchronized_e2e_ms": { + "median": 0.100288 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.7192307692307691, + "submission": 3.309168443496802, + "synchronized_e2e": 1.8618478781110401 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.823304, + "after_init_synchronized_e2e_ms_per_call": 7.850184, + "including_init_host_enqueue_ms_per_call": 44.757742, + "including_init_synchronized_e2e_ms_per_call": 462.96150300000005, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9164232, + "after_init_synchronized_e2e_ms_per_call": 0.9530673000000001, + "including_init_host_enqueue_ms_per_call": 4.6098669999999995, + "including_init_synchronized_e2e_ms_per_call": 46.46419920000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22573512, + "after_init_synchronized_e2e_ms_per_call": 0.26335562999999995, + "including_init_host_enqueue_ms_per_call": 0.5950795, + "including_init_synchronized_e2e_ms_per_call": 4.814468820000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15666631200000003, + "after_init_synchronized_e2e_ms_per_call": 0.194384463, + "including_init_host_enqueue_ms_per_call": 0.19360075000000004, + "including_init_synchronized_e2e_ms_per_call": 0.649495782, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.823304, + "synchronized_e2e_ms": 7.850184, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 9.756202, + "median": 0.085824, + "min": 0.076672, + "p90": 0.10464000000000001, + "sample_count": 2875 + }, + "host_enqueue_ms": { + "max": 44.852847, + "median": 0.148992, + "min": 0.12128, + "p90": 0.20401919999999998, + "sample_count": 2875 + }, + "sample_count": 2875, + "synchronized_e2e_ms": { + "max": 45.149871, + "median": 0.186721, + "min": 0.160832, + "p90": 0.23665920000000001, + "sample_count": 2875 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 2567, + "candidate_precomputed_gpu_span_ms": 0.037088, + "candidate_precomputed_host_enqueue_ms": 0.056384, + "candidate_precomputed_inter_kernel_gap_ms": 0.005664, + "candidate_precomputed_kernel_sum_ms": 0.031392, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.056384, + "candidate_precomputed_synchronized_e2e_ms": 0.072064, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.031392 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.069952, + "synchronized_e2e_ms": 0.087296 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.037088 + }, + "host_enqueue_ms": { + "median": 0.056384 + }, + "inter_kernel_gap_ms": { + "median": 0.005664 + }, + "kernel_sum_ms": { + "median": 0.031392 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2567, + "submission_ms": { + "median": 0.056384 + }, + "synchronized_e2e_ms": { + "median": 0.072064 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7d24f560", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7d24f1d0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 0.9404659188955996, + "submission": 0.8178206583427924, + "synchronized_e2e": 1.1363232682060391 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2875, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.03488, + "candidate_public_raw_host_enqueue_ms": 0.046112, + "candidate_public_raw_inter_kernel_gap_ms": 0.000128, + "candidate_public_raw_kernel_sum_ms": 0.034784, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.046112, + "candidate_public_raw_synchronized_e2e_ms": 0.081888, + "candidate_public_raw_tflops_from_gpu_span": 14.42994495412844, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.034784 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.053824, + "synchronized_e2e_ms": 0.08192 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.03488 + }, + "host_enqueue_ms": { + "median": 0.046112 + }, + "inter_kernel_gap_ms": { + "median": 0.000128 + }, + "kernel_sum_ms": { + "median": 0.034784 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2875, + "submission_ms": { + "median": 0.046112 + }, + "synchronized_e2e_ms": { + "median": 0.081888 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.422979, + "after_init_synchronized_e2e_ms_per_call": 2.448643, + "including_init_host_enqueue_ms_per_call": 39.781001999999994, + "including_init_synchronized_e2e_ms_per_call": 39.884266, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2837987, + "after_init_synchronized_e2e_ms_per_call": 0.3185635, + "including_init_host_enqueue_ms_per_call": 4.019601, + "including_init_synchronized_e2e_ms_per_call": 4.0621258, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06988067, + "after_init_synchronized_e2e_ms_per_call": 0.10555555, + "including_init_host_enqueue_ms_per_call": 0.4434609, + "including_init_synchronized_e2e_ms_per_call": 0.47991177999999995, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.048488867, + "after_init_synchronized_e2e_ms_per_call": 0.084254755, + "including_init_host_enqueue_ms_per_call": 0.08584689, + "including_init_synchronized_e2e_ms_per_call": 0.121690378, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.422979, + "synchronized_e2e_ms": 2.448643, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.03568, + "median": 0.03488, + "min": 0.034624, + "p90": 0.035328, + "sample_count": 2875 + }, + "host_enqueue_ms": { + "max": 1.328482, + "median": 0.046112, + "min": 0.034304, + "p90": 0.0649792, + "sample_count": 2875 + }, + "sample_count": 2875, + "synchronized_e2e_ms": { + "max": 1.558145, + "median": 0.081888, + "min": 0.071456, + "p90": 0.0988358, + "sample_count": 2875 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.066976, + "submission_ms": 0.066976, + "synchronized_e2e_ms": 0.10992 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.045024, + "submission_ms": 0.045024, + "synchronized_e2e_ms": 0.089312 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.19568, + "submission_ms": 0.19568, + "synchronized_e2e_ms": 0.23056 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.823304, + "submission_ms": 7.823304, + "synchronized_e2e_ms": 7.850184 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.068032, + "submission_ms": 0.068032, + "synchronized_e2e_ms": 0.085632 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.069952, + "submission_ms": 0.069952, + "synchronized_e2e_ms": 0.087296 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.22045, + "submission_ms": 1.22045, + "synchronized_e2e_ms": 1.246594 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.129985, + "submission_ms": 1.129985, + "synchronized_e2e_ms": 1.152353 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.053824, + "submission_ms": 0.053824, + "synchronized_e2e_ms": 0.08192 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.422979, + "submission_ms": 2.422979, + "synchronized_e2e_ms": 2.448643 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.916111, + "evolution_kernel_ms": 0.19392, + "evolution_speedup": 4.7242, + "evolution_tflops": 2.5955, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_d9d5_d48_highk_heldout_b2_n640_k4096_d48", + "measurement_order": [ + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 10884, + "measurement_schedule_sha256": "146b227f20f7415c1959f258b2c57e58f5cf832dc1485d0ce33a8ccfefa29ddb", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 2567, + "public_pair_count": 2875, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 2567, + "baseline_public_raw": 2875, + "candidate_precomputed": 2567, + "candidate_public_raw": 2875 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2178 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:adjacent_d9d5_d48_highk_heldout_b2_n640_k4096_d48", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.345987920621225, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.280199785072294, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.2059324286962205, + "including_init_synchronized_e2e_speedup": 11.607622489530083, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.991765534971835, + "including_init_synchronized_e2e_speedup": 11.438394940895236, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.4949482049972733, + "including_init_synchronized_e2e_speedup": 10.031987170642074, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.307103771175882, + "including_init_synchronized_e2e_speedup": 5.337281325562158, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.460550458715596, + "hot_synchronized_e2e_speedup": 2.280199785072294, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 9504801, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_d9d5_d48_highk_heldout_b2_n640_k4096_d48", + "source": "heldout_neighborhood", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 128, + "BLOCK_N": 64, + "D_PAD": 64, + "kernel": "triton_h200_small_d", + "num_stages": 4, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 112, + "K": 1024, + "N": 3712, + "baseline_07cf_adapter_bench_iters": 3666, + "baseline_07cf_adapter_gpu_span_ms": 0.064592, + "baseline_07cf_adapter_host_enqueue_ms": 0.143008, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.029504, + "baseline_07cf_adapter_kernel_sum_ms": 0.035072, + "baseline_07cf_adapter_submission_ms": 0.143008, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.1636965, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.035072 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.295937, + "synchronized_e2e_ms": 0.320769 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.064592 + }, + "host_enqueue_ms": { + "median": 0.143008 + }, + "inter_kernel_gap_ms": { + "median": 0.029504 + }, + "kernel_sum_ms": { + "median": 0.035072 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3666, + "submission_ms": { + "median": 0.143008 + }, + "synchronized_e2e_ms": { + "median": 0.1636965 + } + }, + "baseline_07cf_precomputed_bench_iters": 4703, + "baseline_07cf_precomputed_gpu_span_ms": 0.025824, + "baseline_07cf_precomputed_host_enqueue_ms": 0.040832, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.025824, + "baseline_07cf_precomputed_submission_ms": 0.040832, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.073408, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.025824 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.100608, + "synchronized_e2e_ms": 0.129984 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.025824 + }, + "host_enqueue_ms": { + "median": 0.040832 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.025824 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4703, + "submission_ms": { + "median": 0.040832 + }, + "synchronized_e2e_ms": { + "median": 0.073408 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.501239157372986, + "submission": 3.5023510971786833, + "synchronized_e2e": 2.2299545008718393 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.764873, + "after_init_synchronized_e2e_ms_per_call": 8.798729, + "including_init_host_enqueue_ms_per_call": 42.99278, + "including_init_synchronized_e2e_ms_per_call": 43.144172000000005, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 1.0051945, + "after_init_synchronized_e2e_ms_per_call": 1.0271997499999999, + "including_init_host_enqueue_ms_per_call": 4.4279852, + "including_init_synchronized_e2e_ms_per_call": 4.461744050000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22922664999999998, + "after_init_synchronized_e2e_ms_per_call": 0.25004682500000003, + "including_init_host_enqueue_ms_per_call": 0.57150572, + "including_init_synchronized_e2e_ms_per_call": 0.5935012550000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.151629865, + "after_init_synchronized_e2e_ms_per_call": 0.1723315325, + "including_init_host_enqueue_ms_per_call": 0.185857772, + "including_init_synchronized_e2e_ms_per_call": 0.2066769755, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.764873, + "synchronized_e2e_ms": 8.798729, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.398497, + "median": 0.064592, + "min": 0.047616, + "p90": 0.083152, + "sample_count": 3666 + }, + "host_enqueue_ms": { + "max": 1.50669, + "median": 0.143008, + "min": 0.12, + "p90": 0.18065599999999998, + "sample_count": 3666 + }, + "sample_count": 3666, + "synchronized_e2e_ms": { + "max": 1.620162, + "median": 0.1636965, + "min": 0.141792, + "p90": 0.20385599999999998, + "sample_count": 3666 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 4703, + "candidate_precomputed_gpu_span_ms": 0.01856, + "candidate_precomputed_host_enqueue_ms": 0.053376, + "candidate_precomputed_inter_kernel_gap_ms": 0.003616, + "candidate_precomputed_kernel_sum_ms": 0.014912, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.053376, + "candidate_precomputed_synchronized_e2e_ms": 0.06544, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.014912 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.069088, + "synchronized_e2e_ms": 0.092928 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.01856 + }, + "host_enqueue_ms": { + "median": 0.053376 + }, + "inter_kernel_gap_ms": { + "median": 0.003616 + }, + "kernel_sum_ms": { + "median": 0.014912 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4703, + "submission_ms": { + "median": 0.053376 + }, + "synchronized_e2e_ms": { + "median": 0.06544 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04776000", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc047751f0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.4793103448275862, + "submission": 0.8207434052757794, + "synchronized_e2e": 1.1183374083129585 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 3666, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.027456, + "candidate_public_raw_host_enqueue_ms": 0.043808, + "candidate_public_raw_inter_kernel_gap_ms": 9.6e-05, + "candidate_public_raw_kernel_sum_ms": 0.027328, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.043808, + "candidate_public_raw_synchronized_e2e_ms": 0.073184, + "candidate_public_raw_tflops_from_gpu_span": 124.04482983682983, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.027328 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.076192, + "synchronized_e2e_ms": 0.099136 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.027456 + }, + "host_enqueue_ms": { + "median": 0.043808 + }, + "inter_kernel_gap_ms": { + "median": 9.6e-05 + }, + "kernel_sum_ms": { + "median": 0.027328 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3666, + "submission_ms": { + "median": 0.043808 + }, + "synchronized_e2e_ms": { + "median": 0.073184 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.841763, + "after_init_synchronized_e2e_ms_per_call": 2.895139, + "including_init_host_enqueue_ms_per_call": 37.419303, + "including_init_synchronized_e2e_ms_per_call": 450.98293099999995, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.3236035, + "after_init_synchronized_e2e_ms_per_call": 0.3553795, + "including_init_host_enqueue_ms_per_call": 3.7813575, + "including_init_synchronized_e2e_ms_per_call": 45.164158699999994, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07178755, + "after_init_synchronized_e2e_ms_per_call": 0.10140355, + "including_init_host_enqueue_ms_per_call": 0.41756295, + "including_init_synchronized_e2e_ms_per_call": 4.58228147, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.046605955000000004, + "after_init_synchronized_e2e_ms_per_call": 0.076005955, + "including_init_host_enqueue_ms_per_call": 0.081183495, + "including_init_synchronized_e2e_ms_per_call": 0.5240937469999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.841763, + "synchronized_e2e_ms": 2.895139, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.028064, + "median": 0.027456, + "min": 0.026848, + "p90": 0.02768, + "sample_count": 3666 + }, + "host_enqueue_ms": { + "max": 31.047008, + "median": 0.043808, + "min": 0.0336, + "p90": 0.060176, + "sample_count": 3666 + }, + "sample_count": 3666, + "synchronized_e2e_ms": { + "max": 31.13872, + "median": 0.073184, + "min": 0.06448, + "p90": 0.087392, + "sample_count": 3666 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.10096, + "submission_ms": 0.10096, + "synchronized_e2e_ms": 0.128416 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.100608, + "submission_ms": 0.100608, + "synchronized_e2e_ms": 0.129984 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.295937, + "submission_ms": 0.295937, + "synchronized_e2e_ms": 0.320769 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.764873, + "submission_ms": 8.764873, + "synchronized_e2e_ms": 8.798729 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.100544, + "submission_ms": 0.100544, + "synchronized_e2e_ms": 0.127936 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.069088, + "submission_ms": 0.069088, + "synchronized_e2e_ms": 0.092928 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.519169, + "submission_ms": 1.519169, + "synchronized_e2e_ms": 1.586881 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.512962, + "submission_ms": 1.512962, + "synchronized_e2e_ms": 1.544706 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.076192, + "submission_ms": 0.076192, + "synchronized_e2e_ms": 0.099136 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.841763, + "submission_ms": 2.841763, + "synchronized_e2e_ms": 2.895139 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.386799, + "evolution_kernel_ms": 0.176608, + "evolution_speedup": 2.1902, + "evolution_tflops": 19.2844, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_ef00_d112_odd_tail_b4_n3712_k1024_d112", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 16738, + "measurement_schedule_sha256": "7b069825222cf5c26cadc4028745da3bb4e5c128ce1a0b907a7956d8f0149b19", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 4703, + "public_pair_count": 3666, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 4703, + "baseline_public_raw": 3666, + "candidate_precomputed": 4703, + "candidate_public_raw": 3666 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 3350 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:adjacent_ef00_d112_odd_tail_b4_n3712_k1024_d112", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.3913793103448275, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.2367798972452997, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.039138708020582, + "including_init_synchronized_e2e_speedup": 0.09566697325846243, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.8904305116080127, + "including_init_synchronized_e2e_speedup": 0.09878948658463556, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.465858690351571, + "including_init_synchronized_e2e_speedup": 0.1295209076277019, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.2673425062549373, + "including_init_synchronized_e2e_speedup": 0.3943511569887134, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.352564102564102, + "hot_synchronized_e2e_speedup": 2.2367798972452997, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 9001121, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_ef00_d112_odd_tail_b4_n3712_k1024_d112", + "source": "tail_divisibility", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 128, + "K": 512, + "N": 7552, + "baseline_07cf_adapter_bench_iters": 3054, + "baseline_07cf_adapter_gpu_span_ms": 0.063152, + "baseline_07cf_adapter_host_enqueue_ms": 0.153344, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.025856, + "baseline_07cf_adapter_kernel_sum_ms": 0.037344, + "baseline_07cf_adapter_submission_ms": 0.153344, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.172176, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.037344 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.22864, + "synchronized_e2e_ms": 0.250464 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.063152 + }, + "host_enqueue_ms": { + "median": 0.153344 + }, + "inter_kernel_gap_ms": { + "median": 0.025856 + }, + "kernel_sum_ms": { + "median": 0.037344 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3054, + "submission_ms": { + "median": 0.153344 + }, + "synchronized_e2e_ms": { + "median": 0.172176 + } + }, + "baseline_07cf_precomputed_bench_iters": 7477, + "baseline_07cf_precomputed_gpu_span_ms": 0.0176, + "baseline_07cf_precomputed_host_enqueue_ms": 0.041761, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.0176, + "baseline_07cf_precomputed_submission_ms": 0.041761, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.066688, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.0176 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.055328, + "synchronized_e2e_ms": 0.080448 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.0176 + }, + "host_enqueue_ms": { + "median": 0.041761 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.0176 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 7477, + "submission_ms": { + "median": 0.041761 + }, + "synchronized_e2e_ms": { + "median": 0.066688 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 3.588181818181818, + "submission": 3.671942721678121, + "synchronized_e2e": 2.5818138195777354 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.626729, + "after_init_synchronized_e2e_ms_per_call": 8.657609, + "including_init_host_enqueue_ms_per_call": 43.076173, + "including_init_synchronized_e2e_ms_per_call": 492.699103, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 1.0006825, + "after_init_synchronized_e2e_ms_per_call": 1.0207193, + "including_init_host_enqueue_ms_per_call": 4.4456269, + "including_init_synchronized_e2e_ms_per_call": 49.4248687, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23807785000000004, + "after_init_synchronized_e2e_ms_per_call": 0.25703033000000003, + "including_init_host_enqueue_ms_per_call": 0.58257229, + "including_init_synchronized_e2e_ms_per_call": 5.09744527, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.16181738500000004, + "after_init_synchronized_e2e_ms_per_call": 0.180661433, + "including_init_host_enqueue_ms_per_call": 0.19626682900000003, + "including_init_synchronized_e2e_ms_per_call": 0.664702927, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.626729, + "synchronized_e2e_ms": 8.657609, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.630593, + "median": 0.063152, + "min": 0.052896, + "p90": 0.08292480000000002, + "sample_count": 3054 + }, + "host_enqueue_ms": { + "max": 1.121761, + "median": 0.153344, + "min": 0.129824, + "p90": 0.20921030000000013, + "sample_count": 3054 + }, + "sample_count": 3054, + "synchronized_e2e_ms": { + "max": 1.275873, + "median": 0.172176, + "min": 0.146944, + "p90": 0.23194950000000003, + "sample_count": 3054 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 7477, + "candidate_precomputed_gpu_span_ms": 0.013376, + "candidate_precomputed_host_enqueue_ms": 0.041696, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.013376, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.041696, + "candidate_precomputed_synchronized_e2e_ms": 0.053536, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.013376 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.041088, + "synchronized_e2e_ms": 0.057568 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.013376 + }, + "host_enqueue_ms": { + "median": 0.041696 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.013376 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 7477, + "submission_ms": { + "median": 0.041696 + }, + "synchronized_e2e_ms": { + "median": 0.053536 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc047108f0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc047108c0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.4617224880382773, + "submission": 1.1139677666922487, + "synchronized_e2e": 1.5143548266586968 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 3054, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.032928, + "candidate_public_raw_host_enqueue_ms": 0.046448, + "candidate_public_raw_inter_kernel_gap_ms": 0.00016, + "candidate_public_raw_kernel_sum_ms": 0.032768, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.046448, + "candidate_public_raw_synchronized_e2e_ms": 0.08107249999999999, + "candidate_public_raw_tflops_from_gpu_span": 120.24486686103013, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.032768 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.257665, + "synchronized_e2e_ms": 0.283681 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.032928 + }, + "host_enqueue_ms": { + "median": 0.046448 + }, + "inter_kernel_gap_ms": { + "median": 0.00016 + }, + "kernel_sum_ms": { + "median": 0.032768 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3054, + "submission_ms": { + "median": 0.046448 + }, + "synchronized_e2e_ms": { + "median": 0.08107249999999999 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.214787, + "after_init_synchronized_e2e_ms_per_call": 3.244419, + "including_init_host_enqueue_ms_per_call": 38.088871, + "including_init_synchronized_e2e_ms_per_call": 38.187495, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.3632819, + "after_init_synchronized_e2e_ms_per_call": 0.39740715, + "including_init_host_enqueue_ms_per_call": 3.8506902999999992, + "including_init_synchronized_e2e_ms_per_call": 3.89171475, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07813139, + "after_init_synchronized_e2e_ms_per_call": 0.112705965, + "including_init_host_enqueue_ms_per_call": 0.42687222999999996, + "including_init_synchronized_e2e_ms_per_call": 0.462136725, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.049616339, + "after_init_synchronized_e2e_ms_per_call": 0.08423584649999998, + "including_init_host_enqueue_ms_per_call": 0.084490423, + "including_init_synchronized_e2e_ms_per_call": 0.11917892249999998, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.214787, + "synchronized_e2e_ms": 3.244419, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.033568, + "median": 0.032928, + "min": 0.032352, + "p90": 0.03320670000000001, + "sample_count": 3054 + }, + "host_enqueue_ms": { + "max": 99.782439, + "median": 0.046448, + "min": 0.037728, + "p90": 0.068064, + "sample_count": 3054 + }, + "sample_count": 3054, + "synchronized_e2e_ms": { + "max": 100.123143, + "median": 0.08107249999999999, + "min": 0.073056, + "p90": 0.0996384, + "sample_count": 3054 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.0816, + "submission_ms": 0.0816, + "synchronized_e2e_ms": 0.104096 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.055328, + "submission_ms": 0.055328, + "synchronized_e2e_ms": 0.080448 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.22864, + "submission_ms": 0.22864, + "synchronized_e2e_ms": 0.250464 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.626729, + "submission_ms": 8.626729, + "synchronized_e2e_ms": 8.657609 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.0688, + "submission_ms": 0.0688, + "synchronized_e2e_ms": 0.087616 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.041088, + "submission_ms": 0.041088, + "synchronized_e2e_ms": 0.057568 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.395905, + "submission_ms": 1.395905, + "synchronized_e2e_ms": 1.420417 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.228066, + "submission_ms": 1.228066, + "synchronized_e2e_ms": 1.249346 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.257665, + "submission_ms": 0.257665, + "synchronized_e2e_ms": 0.283681 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.214787, + "submission_ms": 3.214787, + "synchronized_e2e_ms": 3.244419 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.314032, + "evolution_kernel_ms": 0.153023, + "evolution_speedup": 2.0522, + "evolution_tflops": 25.8747, + "expected_route": "aligned_weave_v10_fallback", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_ef00_d128_forced_fallback_b4_n7552_k512_d128", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 21062, + "measurement_schedule_sha256": "d1e969d940f9ed7aacf28d6833f0adfd29e2a96b34c919d5d9305f6d7e0afbfb", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 7477, + "public_pair_count": 3054, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 7477, + "baseline_public_raw": 3054, + "candidate_precomputed": 7477, + "candidate_public_raw": 3054 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 4214 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:adjacent_ef00_d128_forced_fallback_b4_n7552_k512_d128", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.3157894736842106, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.1237287612939038, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.6684620574592866, + "including_init_synchronized_e2e_speedup": 12.902105859522862, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.5684472461051597, + "including_init_synchronized_e2e_speedup": 12.700023479367289, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.280538833947254, + "including_init_synchronized_e2e_speedup": 11.030167035523956, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.144709651608951, + "including_init_synchronized_e2e_speedup": 5.577353050829942, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.9178814382896017, + "hot_synchronized_e2e_speedup": 2.1237287612939038, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 9001281, + "selected_route": "aligned_weave_v10_fallback", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_ef00_d128_forced_fallback_b4_n7552_k512_d128", + "source": "forced_fallback", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 224, + "K": 512, + "N": 5376, + "baseline_07cf_adapter_bench_iters": 3514, + "baseline_07cf_adapter_gpu_span_ms": 0.071904, + "baseline_07cf_adapter_host_enqueue_ms": 0.17648, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.039392, + "baseline_07cf_adapter_kernel_sum_ms": 0.03248, + "baseline_07cf_adapter_submission_ms": 0.17648, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.198112, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.03248 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.235617, + "synchronized_e2e_ms": 0.258401 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.071904 + }, + "host_enqueue_ms": { + "median": 0.17648 + }, + "inter_kernel_gap_ms": { + "median": 0.039392 + }, + "kernel_sum_ms": { + "median": 0.03248 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3514, + "submission_ms": { + "median": 0.17648 + }, + "synchronized_e2e_ms": { + "median": 0.198112 + } + }, + "baseline_07cf_precomputed_bench_iters": 5556, + "baseline_07cf_precomputed_gpu_span_ms": 0.018112, + "baseline_07cf_precomputed_host_enqueue_ms": 0.054048, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.018112, + "baseline_07cf_precomputed_submission_ms": 0.054048, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.078464, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.018112 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.058144, + "synchronized_e2e_ms": 0.080128 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.018112 + }, + "host_enqueue_ms": { + "median": 0.054048 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.018112 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5556, + "submission_ms": { + "median": 0.054048 + }, + "synchronized_e2e_ms": { + "median": 0.078464 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 3.969964664310954, + "submission": 3.265245707519242, + "synchronized_e2e": 2.524877650897227 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.678377, + "after_init_synchronized_e2e_ms_per_call": 8.707145, + "including_init_host_enqueue_ms_per_call": 44.229645, + "including_init_synchronized_e2e_ms_per_call": 44.338381999999996, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 1.0266696999999998, + "after_init_synchronized_e2e_ms_per_call": 1.0490153000000002, + "including_init_host_enqueue_ms_per_call": 4.5817965, + "including_init_synchronized_e2e_ms_per_call": 4.612139, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.26149896999999994, + "after_init_synchronized_e2e_ms_per_call": 0.28320233, + "including_init_host_enqueue_ms_per_call": 0.6170116499999999, + "including_init_synchronized_e2e_ms_per_call": 0.6395147, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.184981897, + "after_init_synchronized_e2e_ms_per_call": 0.206621033, + "including_init_host_enqueue_ms_per_call": 0.220533165, + "including_init_synchronized_e2e_ms_per_call": 0.24225227000000002, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.678377, + "synchronized_e2e_ms": 8.707145, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.328004, + "median": 0.071904, + "min": 0.058784, + "p90": 0.094848, + "sample_count": 3514 + }, + "host_enqueue_ms": { + "max": 59.493181, + "median": 0.17648, + "min": 0.14624, + "p90": 0.2207079, + "sample_count": 3514 + }, + "sample_count": 3514, + "synchronized_e2e_ms": { + "max": 59.674013, + "median": 0.198112, + "min": 0.165728, + "p90": 0.24805790000000003, + "sample_count": 3514 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 5556, + "candidate_precomputed_gpu_span_ms": 0.019488, + "candidate_precomputed_host_enqueue_ms": 0.067392, + "candidate_precomputed_inter_kernel_gap_ms": 0.006464, + "candidate_precomputed_kernel_sum_ms": 0.012992, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.067392, + "candidate_precomputed_synchronized_e2e_ms": 0.080896, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.012992 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.071649, + "synchronized_e2e_ms": 0.089345 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.019488 + }, + "host_enqueue_ms": { + "median": 0.067392 + }, + "inter_kernel_gap_ms": { + "median": 0.006464 + }, + "kernel_sum_ms": { + "median": 0.012992 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5556, + "submission_ms": { + "median": 0.067392 + }, + "synchronized_e2e_ms": { + "median": 0.080896 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb93393cb0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb93390620" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.4548440065681445, + "submission": 0.9204727564102566, + "synchronized_e2e": 1.1396360759493671 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 3514, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.028352, + "candidate_public_raw_host_enqueue_ms": 0.062032500000000004, + "candidate_public_raw_inter_kernel_gap_ms": 3.2e-05, + "candidate_public_raw_kernel_sum_ms": 0.028448, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.062032500000000004, + "candidate_public_raw_synchronized_e2e_ms": 0.092192, + "candidate_public_raw_tflops_from_gpu_span": 86.98683521444696, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.028352 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.085408, + "synchronized_e2e_ms": 0.108768 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.028352 + }, + "host_enqueue_ms": { + "median": 0.062032500000000004 + }, + "inter_kernel_gap_ms": { + "median": 3.2e-05 + }, + "kernel_sum_ms": { + "median": 0.028448 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3514, + "submission_ms": { + "median": 0.062032500000000004 + }, + "synchronized_e2e_ms": { + "median": 0.092192 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.179139, + "after_init_synchronized_e2e_ms_per_call": 3.207555, + "including_init_host_enqueue_ms_per_call": 39.017159, + "including_init_synchronized_e2e_ms_per_call": 487.307768, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.37374315, + "after_init_synchronized_e2e_ms_per_call": 0.40372830000000004, + "including_init_host_enqueue_ms_per_call": 3.95754515, + "including_init_synchronized_e2e_ms_per_call": 48.8137496, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.093203565, + "after_init_synchronized_e2e_ms_per_call": 0.12334563, + "including_init_host_enqueue_ms_per_call": 0.451583765, + "including_init_synchronized_e2e_ms_per_call": 4.96434776, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.0651496065, + "after_init_synchronized_e2e_ms_per_call": 0.09530736299999999, + "including_init_host_enqueue_ms_per_call": 0.10098762650000001, + "including_init_synchronized_e2e_ms_per_call": 0.579407576, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.179139, + "synchronized_e2e_ms": 3.207555, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.028992, + "median": 0.028352, + "min": 0.027936, + "p90": 0.02864, + "sample_count": 3514 + }, + "host_enqueue_ms": { + "max": 74.559085, + "median": 0.062032500000000004, + "min": 0.046272, + "p90": 0.08046080000000001, + "sample_count": 3514 + }, + "sample_count": 3514, + "synchronized_e2e_ms": { + "max": 74.862989, + "median": 0.092192, + "min": 0.077568, + "p90": 0.10911040000000001, + "sample_count": 3514 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.087328, + "submission_ms": 0.087328, + "synchronized_e2e_ms": 0.108704 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.058144, + "submission_ms": 0.058144, + "synchronized_e2e_ms": 0.080128 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.235617, + "submission_ms": 0.235617, + "synchronized_e2e_ms": 0.258401 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.678377, + "submission_ms": 8.678377, + "synchronized_e2e_ms": 8.707145 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.109664, + "submission_ms": 0.109664, + "synchronized_e2e_ms": 0.131296 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.071649, + "submission_ms": 0.071649, + "synchronized_e2e_ms": 0.089345 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.387969, + "submission_ms": 1.387969, + "synchronized_e2e_ms": 1.415041 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.139489, + "submission_ms": 1.139489, + "synchronized_e2e_ms": 1.165537 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.085408, + "submission_ms": 0.085408, + "synchronized_e2e_ms": 0.108768 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.179139, + "submission_ms": 3.179139, + "synchronized_e2e_ms": 3.207555 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.36856, + "evolution_kernel_ms": 0.177808, + "evolution_speedup": 2.0728, + "evolution_tflops": 13.8703, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_ef00_d224_midn_overlap_b2_n5376_k512_d224", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 18140, + "measurement_schedule_sha256": "8e74451389f3198fccc038fa2162023ef42c704cee12122797f421929f22b160", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 5556, + "public_pair_count": 3514, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 5556, + "baseline_public_raw": 3514, + "candidate_precomputed": 5556, + "candidate_public_raw": 3514 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 3630 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:adjacent_ef00_d224_midn_overlap_b2_n5376_k512_d224", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.929392446633826, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.1489066296424855, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.714573873246133, + "including_init_synchronized_e2e_speedup": 0.09098640512539499, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.5983199592399147, + "including_init_synchronized_e2e_speedup": 0.09448442370835614, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.2960061900855346, + "including_init_synchronized_e2e_speedup": 0.12882149497117423, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.1679440758422834, + "including_init_synchronized_e2e_speedup": 0.41810338703614053, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.536117381489842, + "hot_synchronized_e2e_speedup": 2.1489066296424855, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 9002241, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_ef00_d224_midn_overlap_b2_n5376_k512_d224", + "source": "guard_overlap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 288, + "K": 4096, + "N": 1664, + "baseline_07cf_adapter_bench_iters": 2660, + "baseline_07cf_adapter_gpu_span_ms": 0.1523045, + "baseline_07cf_adapter_host_enqueue_ms": 0.16936, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.040768, + "baseline_07cf_adapter_kernel_sum_ms": 0.111552, + "baseline_07cf_adapter_submission_ms": 0.16936, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.265728, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.111552 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.29584, + "synchronized_e2e_ms": 0.381921 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.1523045 + }, + "host_enqueue_ms": { + "median": 0.16936 + }, + "inter_kernel_gap_ms": { + "median": 0.040768 + }, + "kernel_sum_ms": { + "median": 0.111552 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2660, + "submission_ms": { + "median": 0.16936 + }, + "synchronized_e2e_ms": { + "median": 0.265728 + } + }, + "baseline_07cf_precomputed_bench_iters": 3225, + "baseline_07cf_precomputed_gpu_span_ms": 0.115904, + "baseline_07cf_precomputed_host_enqueue_ms": 0.050368, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.115904, + "baseline_07cf_precomputed_submission_ms": 0.050368, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.171648, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.115904 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.065056, + "synchronized_e2e_ms": 0.161569 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.115904 + }, + "host_enqueue_ms": { + "median": 0.050368 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.115904 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3225, + "submission_ms": { + "median": 0.050368 + }, + "synchronized_e2e_ms": { + "median": 0.171648 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.3140573233020432, + "submission": 3.3624523506988564, + "synchronized_e2e": 1.5480984340044743 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.872681, + "after_init_synchronized_e2e_ms_per_call": 8.954154, + "including_init_host_enqueue_ms_per_call": 45.807119, + "including_init_synchronized_e2e_ms_per_call": 464.06547300000005, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 1.0396921000000001, + "after_init_synchronized_e2e_ms_per_call": 1.1345706000000002, + "including_init_host_enqueue_ms_per_call": 4.7331359, + "including_init_synchronized_e2e_ms_per_call": 46.645702500000006, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.25639321000000004, + "after_init_synchronized_e2e_ms_per_call": 0.35261226, + "including_init_host_enqueue_ms_per_call": 0.6257375900000001, + "including_init_synchronized_e2e_ms_per_call": 4.9037254500000005, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.178063321, + "after_init_synchronized_e2e_ms_per_call": 0.27441642600000005, + "including_init_host_enqueue_ms_per_call": 0.214997759, + "including_init_synchronized_e2e_ms_per_call": 0.7295277450000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.872681, + "synchronized_e2e_ms": 8.954154, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.362693, + "median": 0.1523045, + "min": 0.138688, + "p90": 0.17654730000000002, + "sample_count": 2660 + }, + "host_enqueue_ms": { + "max": 44.895694, + "median": 0.16936, + "min": 0.133184, + "p90": 0.23431359999999998, + "sample_count": 2660 + }, + "sample_count": 2660, + "synchronized_e2e_ms": { + "max": 44.967054, + "median": 0.265728, + "min": 0.234144, + "p90": 0.3235497, + "sample_count": 2660 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3225, + "candidate_precomputed_gpu_span_ms": 0.030976, + "candidate_precomputed_host_enqueue_ms": 0.056896, + "candidate_precomputed_inter_kernel_gap_ms": 0.002208, + "candidate_precomputed_kernel_sum_ms": 0.028704, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.056896, + "candidate_precomputed_synchronized_e2e_ms": 0.070848, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.028704 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.060736, + "synchronized_e2e_ms": 0.081344 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.030976 + }, + "host_enqueue_ms": { + "median": 0.056896 + }, + "inter_kernel_gap_ms": { + "median": 0.002208 + }, + "kernel_sum_ms": { + "median": 0.028704 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3225, + "submission_ms": { + "median": 0.056896 + }, + "synchronized_e2e_ms": { + "median": 0.070848 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04d603b0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04d629f0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.215909090909091, + "submission": 0.9702000140607424, + "synchronized_e2e": 1.3130081300813008 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2660, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.037664, + "candidate_public_raw_host_enqueue_ms": 0.0552005, + "candidate_public_raw_inter_kernel_gap_ms": 0.000128, + "candidate_public_raw_kernel_sum_ms": 0.0375525, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.0552005, + "candidate_public_raw_synchronized_e2e_ms": 0.093024, + "candidate_public_raw_tflops_from_gpu_span": 104.23397790994052, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.0375525 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.350144, + "synchronized_e2e_ms": 0.386816 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.037664 + }, + "host_enqueue_ms": { + "median": 0.0552005 + }, + "inter_kernel_gap_ms": { + "median": 0.000128 + }, + "kernel_sum_ms": { + "median": 0.0375525 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2660, + "submission_ms": { + "median": 0.0552005 + }, + "synchronized_e2e_ms": { + "median": 0.093024 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.450499, + "after_init_synchronized_e2e_ms_per_call": 3.489763, + "including_init_host_enqueue_ms_per_call": 40.808521999999996, + "including_init_synchronized_e2e_ms_per_call": 40.925386, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.39473035, + "after_init_synchronized_e2e_ms_per_call": 0.43269789999999997, + "including_init_host_enqueue_ms_per_call": 4.13053265, + "including_init_synchronized_e2e_ms_per_call": 4.1762602, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.089153485, + "after_init_synchronized_e2e_ms_per_call": 0.12699138999999998, + "including_init_host_enqueue_ms_per_call": 0.46273371499999993, + "including_init_synchronized_e2e_ms_per_call": 0.50134762, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.058595798500000004, + "after_init_synchronized_e2e_ms_per_call": 0.09642073899999999, + "including_init_host_enqueue_ms_per_call": 0.09595382150000001, + "including_init_synchronized_e2e_ms_per_call": 0.13385636199999998, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.450499, + "synchronized_e2e_ms": 3.489763, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.039424, + "median": 0.037664, + "min": 0.036384, + "p90": 0.038656, + "sample_count": 2660 + }, + "host_enqueue_ms": { + "max": 18.659891, + "median": 0.0552005, + "min": 0.04176, + "p90": 0.0818912, + "sample_count": 2660 + }, + "sample_count": 2660, + "synchronized_e2e_ms": { + "max": 18.692467, + "median": 0.093024, + "min": 0.081152, + "p90": 0.1156864, + "sample_count": 2660 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.09536, + "submission_ms": 0.09536, + "synchronized_e2e_ms": 0.190912 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.065056, + "submission_ms": 0.065056, + "synchronized_e2e_ms": 0.161569 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.29584, + "submission_ms": 0.29584, + "synchronized_e2e_ms": 0.381921 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.872681, + "submission_ms": 8.872681, + "synchronized_e2e_ms": 8.954154 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.09808, + "submission_ms": 0.09808, + "synchronized_e2e_ms": 0.123008 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.060736, + "submission_ms": 0.060736, + "synchronized_e2e_ms": 0.081344 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.95725, + "submission_ms": 1.95725, + "synchronized_e2e_ms": 1.989378 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.82317, + "submission_ms": 1.82317, + "synchronized_e2e_ms": 1.85709 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.350144, + "submission_ms": 0.350144, + "synchronized_e2e_ms": 0.386816 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.450499, + "submission_ms": 3.450499, + "synchronized_e2e_ms": 3.489763 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 58.88955, + "evolution_kernel_ms": 0.247408, + "evolution_speedup": 238.0261, + "evolution_tflops": 15.868, + "expected_route": "d288_parent_splitk_hybrid_20260629_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_ef00_d288_highk_heldout_b1_n1664_k4096_d288", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 11770, + "measurement_schedule_sha256": "4ec514571307697d0b0243e44ddf6aa62160e16cc8cab7f426a17983e4e0b5de", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3225, + "public_pair_count": 2660, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3225, + "baseline_public_raw": 2660, + "candidate_precomputed": 3225, + "candidate_public_raw": 2660 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2354 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:adjacent_ef00_d288_highk_heldout_b1_n1664_k4096_d288", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 3.7417355371900825, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.85655314757482, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.565834413397128, + "including_init_synchronized_e2e_speedup": 11.339305950590179, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.6220848310102736, + "including_init_synchronized_e2e_speedup": 11.16925197812148, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.7766627328041693, + "including_init_synchronized_e2e_speedup": 9.781088518980106, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.846031142739946, + "including_init_synchronized_e2e_speedup": 5.450078980930321, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 4.043768585386576, + "hot_synchronized_e2e_speedup": 2.85655314757482, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 9002881, + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_ef00_d288_highk_heldout_b1_n1664_k4096_d288", + "source": "heldout_neighborhood", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 3, + "D": 352, + "K": 8192, + "N": 896, + "baseline_07cf_adapter_bench_iters": 717, + "baseline_07cf_adapter_gpu_span_ms": 0.252288, + "baseline_07cf_adapter_host_enqueue_ms": 0.153536, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.004096, + "baseline_07cf_adapter_kernel_sum_ms": 0.24816, + "baseline_07cf_adapter_submission_ms": 0.153536, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.357921, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.24816 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.189856, + "synchronized_e2e_ms": 0.401696 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.252288 + }, + "host_enqueue_ms": { + "median": 0.153536 + }, + "inter_kernel_gap_ms": { + "median": 0.004096 + }, + "kernel_sum_ms": { + "median": 0.24816 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 717, + "submission_ms": { + "median": 0.153536 + }, + "synchronized_e2e_ms": { + "median": 0.357921 + } + }, + "baseline_07cf_precomputed_bench_iters": 967, + "baseline_07cf_precomputed_gpu_span_ms": 0.247233, + "baseline_07cf_precomputed_host_enqueue_ms": 0.044832, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.247233, + "baseline_07cf_precomputed_submission_ms": 0.044832, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.298881, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.247233 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.044672, + "synchronized_e2e_ms": 0.234113 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.247233 + }, + "host_enqueue_ms": { + "median": 0.044832 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.247233 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 967, + "submission_ms": { + "median": 0.044832 + }, + "synchronized_e2e_ms": { + "median": 0.298881 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.020446299644465, + "submission": 3.4246966452533907, + "synchronized_e2e": 1.197536812309916 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.625641, + "after_init_synchronized_e2e_ms_per_call": 8.811241, + "including_init_host_enqueue_ms_per_call": 42.853548, + "including_init_synchronized_e2e_ms_per_call": 43.156684000000006, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 1.0007465, + "after_init_synchronized_e2e_ms_per_call": 1.2032530000000001, + "including_init_host_enqueue_ms_per_call": 4.4235372, + "including_init_synchronized_e2e_ms_per_call": 4.637797300000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23825705, + "after_init_synchronized_e2e_ms_per_call": 0.4424542, + "including_init_host_enqueue_ms_per_call": 0.58053612, + "including_init_synchronized_e2e_ms_per_call": 0.7859086300000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.162008105, + "after_init_synchronized_e2e_ms_per_call": 0.36637432000000003, + "including_init_host_enqueue_ms_per_call": 0.19623601200000002, + "including_init_synchronized_e2e_ms_per_call": 0.400719763, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.625641, + "synchronized_e2e_ms": 8.811241, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.534081, + "median": 0.252288, + "min": 0.248961, + "p90": 0.26061500000000004, + "sample_count": 717 + }, + "host_enqueue_ms": { + "max": 0.44512, + "median": 0.153536, + "min": 0.131008, + "p90": 0.1866436, + "sample_count": 717 + }, + "sample_count": 717, + "synchronized_e2e_ms": { + "max": 0.634816, + "median": 0.357921, + "min": 0.341281, + "p90": 0.38540220000000003, + "sample_count": 717 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 967, + "candidate_precomputed_gpu_span_ms": 0.104064, + "candidate_precomputed_host_enqueue_ms": 0.053664, + "candidate_precomputed_inter_kernel_gap_ms": 0.002112, + "candidate_precomputed_kernel_sum_ms": 0.101856, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.053664, + "candidate_precomputed_synchronized_e2e_ms": 0.140961, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.101856 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.063904, + "synchronized_e2e_ms": 0.127712 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.104064 + }, + "host_enqueue_ms": { + "median": 0.053664 + }, + "inter_kernel_gap_ms": { + "median": 0.002112 + }, + "kernel_sum_ms": { + "median": 0.101856 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 967, + "submission_ms": { + "median": 0.053664 + }, + "synchronized_e2e_ms": { + "median": 0.140961 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282c4a70", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282c66c0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.3379458794587944, + "submission": 0.9660107334525938, + "synchronized_e2e": 1.3670731620802916 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 717, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.139232, + "candidate_public_raw_host_enqueue_ms": 0.05184, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.139296, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.05184, + "candidate_public_raw_synchronized_e2e_ms": 0.192704, + "candidate_public_raw_tflops_from_gpu_span": 111.34040726269824, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.1392 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.2384, + "synchronized_e2e_ms": 0.343808 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.139232 + }, + "host_enqueue_ms": { + "median": 0.05184 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.139296 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 717, + "submission_ms": { + "median": 0.05184 + }, + "synchronized_e2e_ms": { + "median": 0.192704 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.538946, + "after_init_synchronized_e2e_ms_per_call": 2.636195, + "including_init_host_enqueue_ms_per_call": 37.116486, + "including_init_synchronized_e2e_ms_per_call": 450.72398699999997, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.3005506, + "after_init_synchronized_e2e_ms_per_call": 0.43705309999999997, + "including_init_host_enqueue_ms_per_call": 3.7583046, + "including_init_synchronized_e2e_ms_per_call": 45.2458323, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07671106, + "after_init_synchronized_e2e_ms_per_call": 0.21713891, + "including_init_host_enqueue_ms_per_call": 0.42248646, + "including_init_synchronized_e2e_ms_per_call": 4.698016829999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.054327106, + "after_init_synchronized_e2e_ms_per_call": 0.19514749099999998, + "including_init_host_enqueue_ms_per_call": 0.088904646, + "including_init_synchronized_e2e_ms_per_call": 0.643235283, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.538946, + "synchronized_e2e_ms": 2.636195, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.142592, + "median": 0.139232, + "min": 0.137536, + "p90": 0.14032, + "sample_count": 717 + }, + "host_enqueue_ms": { + "max": 0.1264, + "median": 0.05184, + "min": 0.041024, + "p90": 0.0659776, + "sample_count": 717 + }, + "sample_count": 717, + "synchronized_e2e_ms": { + "max": 0.265984, + "median": 0.192704, + "min": 0.182432, + "p90": 0.2037312, + "sample_count": 717 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.066368, + "submission_ms": 0.066368, + "synchronized_e2e_ms": 0.254464 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.044672, + "submission_ms": 0.044672, + "synchronized_e2e_ms": 0.234113 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.189856, + "submission_ms": 0.189856, + "synchronized_e2e_ms": 0.401696 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.625641, + "submission_ms": 8.625641, + "synchronized_e2e_ms": 8.811241 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.06176, + "submission_ms": 0.06176, + "synchronized_e2e_ms": 0.120032 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.063904, + "submission_ms": 0.063904, + "synchronized_e2e_ms": 0.127712 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.268226, + "submission_ms": 1.268226, + "synchronized_e2e_ms": 1.288834 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.166817, + "submission_ms": 1.166817, + "synchronized_e2e_ms": 1.185217 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.2384, + "submission_ms": 0.2384, + "synchronized_e2e_ms": 0.343808 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.538946, + "submission_ms": 2.538946, + "synchronized_e2e_ms": 2.636195 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 3.136652, + "evolution_kernel_ms": 0.316191, + "evolution_speedup": 9.9201, + "evolution_tflops": 49.0278, + "expected_route": "d352_exactd_splitk_c95c_v2", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_ef00_d352_request_low_n_b3_n896_k8192_d352", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 3368, + "measurement_schedule_sha256": "ca9c7f82e42b63f3096e7ed51c2ffca5e252cea0ead2ceb524d5d964bd4e459d", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 967, + "public_pair_count": 717, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 967, + "baseline_public_raw": 717, + "candidate_precomputed": 967, + "candidate_public_raw": 717 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 676 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:adjacent_ef00_d352_request_low_n_b3_n896_k8192_d352", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 2.3757783671586714, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.857361549319163, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.3424086609677968, + "including_init_synchronized_e2e_speedup": 0.09574969436893983, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.7531048286810007, + "including_init_synchronized_e2e_speedup": 0.1025021988599821, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.0376550660588655, + "including_init_synchronized_e2e_speedup": 0.1672851882908219, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.8774226515676806, + "including_init_synchronized_e2e_speedup": 0.6229754082068909, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.8119972420133303, + "hot_synchronized_e2e_speedup": 1.857361549319163, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 9003521, + "selected_route": "d352_exactd_splitk_c95c_v2", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_ef00_d352_request_low_n_b3_n896_k8192_d352", + "source": "request_specific", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 416, + "K": 768, + "N": 2176, + "baseline_07cf_adapter_bench_iters": 3005, + "baseline_07cf_adapter_gpu_span_ms": 0.068576, + "baseline_07cf_adapter_host_enqueue_ms": 0.148961, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.038176, + "baseline_07cf_adapter_kernel_sum_ms": 0.0304, + "baseline_07cf_adapter_submission_ms": 0.148961, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.171456, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.0304 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.217088, + "synchronized_e2e_ms": 0.238496 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.068576 + }, + "host_enqueue_ms": { + "median": 0.148961 + }, + "inter_kernel_gap_ms": { + "median": 0.038176 + }, + "kernel_sum_ms": { + "median": 0.0304 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3005, + "submission_ms": { + "median": 0.148961 + }, + "synchronized_e2e_ms": { + "median": 0.171456 + } + }, + "baseline_07cf_precomputed_bench_iters": 3714, + "baseline_07cf_precomputed_gpu_span_ms": 0.026912, + "baseline_07cf_precomputed_host_enqueue_ms": 0.043552, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.026912, + "baseline_07cf_precomputed_submission_ms": 0.043552, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.0777925, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.026912 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.044416, + "synchronized_e2e_ms": 0.072448 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.026912 + }, + "host_enqueue_ms": { + "median": 0.043552 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.026912 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3714, + "submission_ms": { + "median": 0.043552 + }, + "synchronized_e2e_ms": { + "median": 0.0777925 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.5481569560047563, + "submission": 3.4203021675238796, + "synchronized_e2e": 2.2040170967638266 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 9.01681, + "after_init_synchronized_e2e_ms_per_call": 9.050794, + "including_init_host_enqueue_ms_per_call": 43.466254, + "including_init_synchronized_e2e_ms_per_call": 493.092288, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 1.0357458999999998, + "after_init_synchronized_e2e_ms_per_call": 1.0593898, + "including_init_host_enqueue_ms_per_call": 4.4806903, + "including_init_synchronized_e2e_ms_per_call": 49.4635392, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23763949, + "after_init_synchronized_e2e_ms_per_call": 0.26024938, + "including_init_host_enqueue_ms_per_call": 0.5821339299999999, + "including_init_synchronized_e2e_ms_per_call": 5.10066432, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.157828849, + "after_init_synchronized_e2e_ms_per_call": 0.180335338, + "including_init_host_enqueue_ms_per_call": 0.192278293, + "including_init_synchronized_e2e_ms_per_call": 0.6643768320000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 9.01681, + "synchronized_e2e_ms": 9.050794, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.781793, + "median": 0.068576, + "min": 0.06, + "p90": 0.08733439999999998, + "sample_count": 3005 + }, + "host_enqueue_ms": { + "max": 79.290322, + "median": 0.148961, + "min": 0.125184, + "p90": 0.20250240000000003, + "sample_count": 3005 + }, + "sample_count": 3005, + "synchronized_e2e_ms": { + "max": 79.55317, + "median": 0.171456, + "min": 0.148481, + "p90": 0.223584, + "sample_count": 3005 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3714, + "candidate_precomputed_gpu_span_ms": 0.032256, + "candidate_precomputed_host_enqueue_ms": 0.057792, + "candidate_precomputed_inter_kernel_gap_ms": 0.006176, + "candidate_precomputed_kernel_sum_ms": 0.026048, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.057792, + "candidate_precomputed_synchronized_e2e_ms": 0.0708645, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.026048 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.047872, + "synchronized_e2e_ms": 0.063584 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.032256 + }, + "host_enqueue_ms": { + "median": 0.057792 + }, + "inter_kernel_gap_ms": { + "median": 0.006176 + }, + "kernel_sum_ms": { + "median": 0.026048 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3714, + "submission_ms": { + "median": 0.057792 + }, + "synchronized_e2e_ms": { + "median": 0.0708645 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01bbf170", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01bbe480" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.0327380952380953, + "submission": 0.7873754152823921, + "synchronized_e2e": 1.1460745507270917 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 3005, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.033312, + "candidate_public_raw_host_enqueue_ms": 0.045504, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.033248, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.045504, + "candidate_public_raw_synchronized_e2e_ms": 0.081216, + "candidate_public_raw_tflops_from_gpu_span": 41.739066282420744, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.033248 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.094752, + "synchronized_e2e_ms": 0.11376 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.033312 + }, + "host_enqueue_ms": { + "median": 0.045504 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.033248 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3005, + "submission_ms": { + "median": 0.045504 + }, + "synchronized_e2e_ms": { + "median": 0.081216 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.614724, + "after_init_synchronized_e2e_ms_per_call": 3.643556, + "including_init_host_enqueue_ms_per_call": 38.488808, + "including_init_synchronized_e2e_ms_per_call": 38.586631999999994, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.402426, + "after_init_synchronized_e2e_ms_per_call": 0.43744999999999995, + "including_init_host_enqueue_ms_per_call": 3.8898344000000002, + "including_init_synchronized_e2e_ms_per_call": 3.9317575999999996, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.08119620000000001, + "after_init_synchronized_e2e_ms_per_call": 0.1168394, + "including_init_host_enqueue_ms_per_call": 0.42993704, + "including_init_synchronized_e2e_ms_per_call": 0.46627016, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.04907322000000001, + "after_init_synchronized_e2e_ms_per_call": 0.08477834, + "including_init_host_enqueue_ms_per_call": 0.083947304, + "including_init_synchronized_e2e_ms_per_call": 0.119721416, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.614724, + "synchronized_e2e_ms": 3.643556, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.033889, + "median": 0.033312, + "min": 0.032896, + "p90": 0.033536, + "sample_count": 3005 + }, + "host_enqueue_ms": { + "max": 15.721744, + "median": 0.045504, + "min": 0.035296, + "p90": 0.06405759999999999, + "sample_count": 3005 + }, + "sample_count": 3005, + "synchronized_e2e_ms": { + "max": 20.750645, + "median": 0.081216, + "min": 0.072448, + "p90": 0.09719039999999998, + "sample_count": 3005 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.073856, + "submission_ms": 0.073856, + "synchronized_e2e_ms": 0.1008 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.044416, + "submission_ms": 0.044416, + "synchronized_e2e_ms": 0.072448 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.217088, + "submission_ms": 0.217088, + "synchronized_e2e_ms": 0.238496 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 9.01681, + "submission_ms": 9.01681, + "synchronized_e2e_ms": 9.050794 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.112, + "submission_ms": 0.112, + "synchronized_e2e_ms": 0.131584 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.047872, + "submission_ms": 0.047872, + "synchronized_e2e_ms": 0.063584 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.695458, + "submission_ms": 1.695458, + "synchronized_e2e_ms": 1.723938 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.231489, + "submission_ms": 1.231489, + "synchronized_e2e_ms": 1.253985 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.094752, + "submission_ms": 0.094752, + "synchronized_e2e_ms": 0.11376 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.614724, + "submission_ms": 3.614724, + "synchronized_e2e_ms": 3.643556 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.529311, + "evolution_kernel_ms": 0.191648, + "evolution_speedup": 2.7619, + "evolution_tflops": 7.255, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_ef00_d416_random_midk_b1_n2176_k768_d416", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 13438, + "measurement_schedule_sha256": "535aed2bf528a961b8d8c1c5d4148db49b14fb4cc5da44cb1863de4faff8723b", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3714, + "public_pair_count": 3005, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3714, + "baseline_public_raw": 3005, + "candidate_precomputed": 3714, + "candidate_public_raw": 3005 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2688 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:adjacent_ef00_d416_random_midk_b1_n2176_k768_d416", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.8343253968253967, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.111111111111111, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.4840551373438475, + "including_init_synchronized_e2e_speedup": 12.778837189003696, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.4217391701908793, + "including_init_synchronized_e2e_speedup": 12.580515950423802, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.22741113014959, + "including_init_synchronized_e2e_speedup": 10.939289616989429, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.1271392905310487, + "including_init_synchronized_e2e_speedup": 5.5493566163634425, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.058597502401537, + "hot_synchronized_e2e_speedup": 2.111111111111111, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 9004161, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_ef00_d416_random_midk_b1_n2176_k768_d416", + "source": "random_legal", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 3, + "D": 480, + "K": 256, + "N": 3200, + "baseline_07cf_adapter_bench_iters": 2886, + "baseline_07cf_adapter_gpu_span_ms": 0.068032, + "baseline_07cf_adapter_host_enqueue_ms": 0.181168, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.038464, + "baseline_07cf_adapter_kernel_sum_ms": 0.029536, + "baseline_07cf_adapter_submission_ms": 0.181168, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.203072, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.029536 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.22064, + "synchronized_e2e_ms": 0.24368 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.068032 + }, + "host_enqueue_ms": { + "median": 0.181168 + }, + "inter_kernel_gap_ms": { + "median": 0.038464 + }, + "kernel_sum_ms": { + "median": 0.029536 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2886, + "submission_ms": { + "median": 0.181168 + }, + "synchronized_e2e_ms": { + "median": 0.203072 + } + }, + "baseline_07cf_precomputed_bench_iters": 6649, + "baseline_07cf_precomputed_gpu_span_ms": 0.014976, + "baseline_07cf_precomputed_host_enqueue_ms": 0.053536, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.014976, + "baseline_07cf_precomputed_submission_ms": 0.053536, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.075072, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.014976 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.056224, + "synchronized_e2e_ms": 0.074464 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.014976 + }, + "host_enqueue_ms": { + "median": 0.053536 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.014976 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 6649, + "submission_ms": { + "median": 0.053536 + }, + "synchronized_e2e_ms": { + "median": 0.075072 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 4.542735042735043, + "submission": 3.3840406455469214, + "synchronized_e2e": 2.7050298380221656 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.459368, + "after_init_synchronized_e2e_ms_per_call": 8.487944, + "including_init_host_enqueue_ms_per_call": 44.010636, + "including_init_synchronized_e2e_ms_per_call": 44.119181, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 1.008988, + "after_init_synchronized_e2e_ms_per_call": 1.0315592, + "including_init_host_enqueue_ms_per_call": 4.5641148000000005, + "including_init_synchronized_e2e_ms_per_call": 4.5946829, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.26394999999999996, + "after_init_synchronized_e2e_ms_per_call": 0.28592072, + "including_init_host_enqueue_ms_per_call": 0.6194626799999999, + "including_init_synchronized_e2e_ms_per_call": 0.64223309, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.1894462, + "after_init_synchronized_e2e_ms_per_call": 0.211356872, + "including_init_host_enqueue_ms_per_call": 0.224997468, + "including_init_synchronized_e2e_ms_per_call": 0.246988109, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.459368, + "synchronized_e2e_ms": 8.487944, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 7.538984, + "median": 0.068032, + "min": 0.05504, + "p90": 0.09236849999999999, + "sample_count": 2886 + }, + "host_enqueue_ms": { + "max": 81.278837, + "median": 0.181168, + "min": 0.146784, + "p90": 0.233136, + "sample_count": 2886 + }, + "sample_count": 2886, + "synchronized_e2e_ms": { + "max": 81.497237, + "median": 0.203072, + "min": 0.16544, + "p90": 0.2606725, + "sample_count": 2886 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 6649, + "candidate_precomputed_gpu_span_ms": 0.023264, + "candidate_precomputed_host_enqueue_ms": 0.06784, + "candidate_precomputed_inter_kernel_gap_ms": 0.005536, + "candidate_precomputed_kernel_sum_ms": 0.017728, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.06784, + "candidate_precomputed_synchronized_e2e_ms": 0.082049, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.017728 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.056576, + "synchronized_e2e_ms": 0.0728 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.023264 + }, + "host_enqueue_ms": { + "median": 0.06784 + }, + "inter_kernel_gap_ms": { + "median": 0.005536 + }, + "kernel_sum_ms": { + "median": 0.017728 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 6649, + "submission_ms": { + "median": 0.06784 + }, + "synchronized_e2e_ms": { + "median": 0.082049 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282c75c0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282c7080" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.5048143053645116, + "submission": 0.9113207547169812, + "synchronized_e2e": 1.1985033333739596 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2886, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.035008, + "candidate_public_raw_host_enqueue_ms": 0.061824, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.035232, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.061824, + "candidate_public_raw_synchronized_e2e_ms": 0.098336, + "candidate_public_raw_tflops_from_gpu_span": 67.39305301645338, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.035008 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.085409, + "synchronized_e2e_ms": 0.107361 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.035008 + }, + "host_enqueue_ms": { + "median": 0.061824 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.035232 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2886, + "submission_ms": { + "median": 0.061824 + }, + "synchronized_e2e_ms": { + "median": 0.098336 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.52874, + "after_init_synchronized_e2e_ms_per_call": 3.556228, + "including_init_host_enqueue_ms_per_call": 39.36676, + "including_init_synchronized_e2e_ms_per_call": 487.656441, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.4085156, + "after_init_synchronized_e2e_ms_per_call": 0.44412520000000005, + "including_init_host_enqueue_ms_per_call": 3.9923176, + "including_init_synchronized_e2e_ms_per_call": 48.8541465, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.09649316, + "after_init_synchronized_e2e_ms_per_call": 0.13291492000000002, + "including_init_host_enqueue_ms_per_call": 0.45487336, + "including_init_synchronized_e2e_ms_per_call": 4.97391705, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06529091599999999, + "after_init_synchronized_e2e_ms_per_call": 0.10179389200000001, + "including_init_host_enqueue_ms_per_call": 0.101128936, + "including_init_synchronized_e2e_ms_per_call": 0.5858941049999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.52874, + "synchronized_e2e_ms": 3.556228, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.035808, + "median": 0.035008, + "min": 0.034176, + "p90": 0.035264, + "sample_count": 2886 + }, + "host_enqueue_ms": { + "max": 54.674616, + "median": 0.061824, + "min": 0.049088, + "p90": 0.082272, + "sample_count": 2886 + }, + "sample_count": 2886, + "synchronized_e2e_ms": { + "max": 54.848601, + "median": 0.098336, + "min": 0.08672, + "p90": 0.11607999999999999, + "sample_count": 2886 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.081696, + "submission_ms": 0.081696, + "synchronized_e2e_ms": 0.100832 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.056224, + "submission_ms": 0.056224, + "synchronized_e2e_ms": 0.074464 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.22064, + "submission_ms": 0.22064, + "synchronized_e2e_ms": 0.24368 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.459368, + "submission_ms": 8.459368, + "synchronized_e2e_ms": 8.487944 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.102561, + "submission_ms": 0.102561, + "synchronized_e2e_ms": 0.121921 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.056576, + "submission_ms": 0.056576, + "synchronized_e2e_ms": 0.0728 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.341729, + "submission_ms": 1.341729, + "synchronized_e2e_ms": 1.367681 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.137345, + "submission_ms": 1.137345, + "synchronized_e2e_ms": 1.162145 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.085409, + "submission_ms": 0.085409, + "synchronized_e2e_ms": 0.107361 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.52874, + "submission_ms": 3.52874, + "synchronized_e2e_ms": 3.556228 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.334224, + "evolution_kernel_ms": 0.18192, + "evolution_speedup": 1.8372, + "evolution_tflops": 12.9689, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_ef00_d480_lowk_boundary_b3_n3200_k256_d480", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 19070, + "measurement_schedule_sha256": "a9638aeddab18baf2ccdaf2029730db86d77c835b852c48498fb640860eba5a7", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 6649, + "public_pair_count": 2886, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 6649, + "baseline_public_raw": 2886, + "candidate_precomputed": 6649, + "candidate_public_raw": 2886 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 3816 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:adjacent_ef00_d480_lowk_boundary_b3_n3200_k256_d480", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.6437414030261348, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.0650829808005207, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.3867828496935517, + "including_init_synchronized_e2e_speedup": 0.09047185126792984, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.3226765785863983, + "including_init_synchronized_e2e_speedup": 0.09404898517672394, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.151155942463043, + "including_init_synchronized_e2e_speedup": 0.12912018506621456, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.076321750228393, + "including_init_synchronized_e2e_speedup": 0.42155759358596046, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.943327239488117, + "hot_synchronized_e2e_speedup": 2.0650829808005207, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 9004801, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_ef00_d480_lowk_boundary_b3_n3200_k256_d480", + "source": "guard_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 3, + "D": 48, + "K": 512, + "N": 3456, + "baseline_07cf_adapter_bench_iters": 5947, + "baseline_07cf_adapter_gpu_span_ms": 0.059488, + "baseline_07cf_adapter_host_enqueue_ms": 0.166688, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.041024, + "baseline_07cf_adapter_kernel_sum_ms": 0.018464, + "baseline_07cf_adapter_submission_ms": 0.166688, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.189376, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.018464 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.236704, + "synchronized_e2e_ms": 0.282848 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.059488 + }, + "host_enqueue_ms": { + "median": 0.166688 + }, + "inter_kernel_gap_ms": { + "median": 0.041024 + }, + "kernel_sum_ms": { + "median": 0.018464 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 5947, + "submission_ms": { + "median": 0.166688 + }, + "synchronized_e2e_ms": { + "median": 0.189376 + } + }, + "baseline_07cf_precomputed_bench_iters": 8669, + "baseline_07cf_precomputed_gpu_span_ms": 0.011584, + "baseline_07cf_precomputed_host_enqueue_ms": 0.04784, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.011584, + "baseline_07cf_precomputed_submission_ms": 0.04784, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.064704, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.011584 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.054688, + "synchronized_e2e_ms": 0.074112 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.011584 + }, + "host_enqueue_ms": { + "median": 0.04784 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.011584 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 8669, + "submission_ms": { + "median": 0.04784 + }, + "synchronized_e2e_ms": { + "median": 0.064704 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 5.135359116022099, + "submission": 3.4842809364548497, + "synchronized_e2e": 2.926805143422354 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.989224, + "after_init_synchronized_e2e_ms_per_call": 8.026088, + "including_init_host_enqueue_ms_per_call": 44.923662, + "including_init_synchronized_e2e_ms_per_call": 463.13740700000005, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9489416, + "after_init_synchronized_e2e_ms_per_call": 0.9730471999999999, + "including_init_host_enqueue_ms_per_call": 4.6423854, + "including_init_synchronized_e2e_ms_per_call": 46.484179100000006, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.24491336, + "after_init_synchronized_e2e_ms_per_call": 0.26774312, + "including_init_host_enqueue_ms_per_call": 0.61425774, + "including_init_synchronized_e2e_ms_per_call": 4.81885631, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.174510536, + "after_init_synchronized_e2e_ms_per_call": 0.19721271199999998, + "including_init_host_enqueue_ms_per_call": 0.211444974, + "including_init_synchronized_e2e_ms_per_call": 0.6523240310000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.989224, + "synchronized_e2e_ms": 8.026088, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.285542, + "median": 0.059488, + "min": 0.039232, + "p90": 0.08098560000000003, + "sample_count": 5947 + }, + "host_enqueue_ms": { + "max": 74.363053, + "median": 0.166688, + "min": 0.124992, + "p90": 0.2271814000000001, + "sample_count": 5947 + }, + "sample_count": 5947, + "synchronized_e2e_ms": { + "max": 74.594605, + "median": 0.189376, + "min": 0.143264, + "p90": 0.25572480000000003, + "sample_count": 5947 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 8669, + "candidate_precomputed_gpu_span_ms": 0.015136, + "candidate_precomputed_host_enqueue_ms": 0.05808, + "candidate_precomputed_inter_kernel_gap_ms": 0.005888, + "candidate_precomputed_kernel_sum_ms": 0.00928, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.05808, + "candidate_precomputed_synchronized_e2e_ms": 0.071616, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.00928 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.059872, + "synchronized_e2e_ms": 0.081408 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.015136 + }, + "host_enqueue_ms": { + "median": 0.05808 + }, + "inter_kernel_gap_ms": { + "median": 0.005888 + }, + "kernel_sum_ms": { + "median": 0.00928 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 8669, + "submission_ms": { + "median": 0.05808 + }, + "synchronized_e2e_ms": { + "median": 0.071616 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc03ca0350", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc03ca0080" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.1331923890063424, + "submission": 0.8699724517906336, + "synchronized_e2e": 0.9490616621983914 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 5947, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.017152, + "candidate_public_raw_host_enqueue_ms": 0.050528, + "candidate_public_raw_inter_kernel_gap_ms": 0.000128, + "candidate_public_raw_kernel_sum_ms": 0.017024, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.050528, + "candidate_public_raw_synchronized_e2e_ms": 0.067968, + "candidate_public_raw_tflops_from_gpu_span": 29.71128358208955, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.017024 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.096608, + "synchronized_e2e_ms": 0.122976 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.017152 + }, + "host_enqueue_ms": { + "median": 0.050528 + }, + "inter_kernel_gap_ms": { + "median": 0.000128 + }, + "kernel_sum_ms": { + "median": 0.017024 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 5947, + "submission_ms": { + "median": 0.050528 + }, + "synchronized_e2e_ms": { + "median": 0.067968 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.791331, + "after_init_synchronized_e2e_ms_per_call": 2.824291, + "including_init_host_enqueue_ms_per_call": 40.149353999999995, + "including_init_synchronized_e2e_ms_per_call": 40.259914, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.3246083, + "after_init_synchronized_e2e_ms_per_call": 0.3436003, + "including_init_host_enqueue_ms_per_call": 4.060410599999999, + "including_init_synchronized_e2e_ms_per_call": 4.0871626, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07793603, + "after_init_synchronized_e2e_ms_per_call": 0.09553123, + "including_init_host_enqueue_ms_per_call": 0.45151625999999995, + "including_init_synchronized_e2e_ms_per_call": 0.46988746, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.053268803, + "after_init_synchronized_e2e_ms_per_call": 0.07072432299999999, + "including_init_host_enqueue_ms_per_call": 0.090626826, + "including_init_synchronized_e2e_ms_per_call": 0.10815994599999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.791331, + "synchronized_e2e_ms": 2.824291, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.017792, + "median": 0.017152, + "min": 0.016672, + "p90": 0.017472, + "sample_count": 5947 + }, + "host_enqueue_ms": { + "max": 83.621782, + "median": 0.050528, + "min": 0.036672, + "p90": 0.071776, + "sample_count": 5947 + }, + "sample_count": 5947, + "synchronized_e2e_ms": { + "max": 83.929111, + "median": 0.067968, + "min": 0.055488, + "p90": 0.09303680000000002, + "sample_count": 5947 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.07984, + "submission_ms": 0.07984, + "synchronized_e2e_ms": 0.100704 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.054688, + "submission_ms": 0.054688, + "synchronized_e2e_ms": 0.074112 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.236704, + "submission_ms": 0.236704, + "synchronized_e2e_ms": 0.282848 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.989224, + "submission_ms": 7.989224, + "synchronized_e2e_ms": 8.026088 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.088032, + "submission_ms": 0.088032, + "synchronized_e2e_ms": 0.11024 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.059872, + "submission_ms": 0.059872, + "synchronized_e2e_ms": 0.081408 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.604002, + "submission_ms": 1.604002, + "synchronized_e2e_ms": 1.63309 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.377025, + "submission_ms": 1.377025, + "synchronized_e2e_ms": 1.404705 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.096608, + "submission_ms": 0.096608, + "synchronized_e2e_ms": 0.122976 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.791331, + "submission_ms": 2.791331, + "synchronized_e2e_ms": 2.824291 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.257807, + "evolution_kernel_ms": 0.172495, + "evolution_speedup": 1.4946, + "evolution_tflops": 2.9543, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "adjacent_ef00_d48_midn_boundary_b3_n3456_k512_d48", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 29232, + "measurement_schedule_sha256": "91cf3d95a653ac5646d1fbcfd073538ac188d248dfcf4890b8944ea9ca9b99f4", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 8669, + "public_pair_count": 5947, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 8669, + "baseline_public_raw": 5947, + "candidate_precomputed": 8669, + "candidate_public_raw": 5947 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 5848 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:adjacent_ef00_d48_midn_boundary_b3_n3456_k512_d48", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.7653276955602537, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.786252354048964, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.8418063152840833, + "including_init_synchronized_e2e_speedup": 11.503685949254637, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.8319160373259273, + "including_init_synchronized_e2e_speedup": 11.373215026972503, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.8026763603902096, + "including_init_synchronized_e2e_speedup": 10.255341374719812, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.7884708348498437, + "including_init_synchronized_e2e_speedup": 6.031105368710152, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 3.468283582089552, + "hot_synchronized_e2e_speedup": 2.786252354048964, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 9000481, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "adjacent_ef00_d48_midn_boundary_b3_n3456_k512_d48", + "source": "guard_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 64, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 112, + "K": 256, + "N": 256, + "baseline_07cf_adapter_bench_iters": 9859, + "baseline_07cf_adapter_gpu_span_ms": 0.049792, + "baseline_07cf_adapter_host_enqueue_ms": 0.145056, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.040864, + "baseline_07cf_adapter_kernel_sum_ms": 0.008928, + "baseline_07cf_adapter_submission_ms": 0.145056, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.165665, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.008928 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.210017, + "synchronized_e2e_ms": 0.257057 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.049792 + }, + "host_enqueue_ms": { + "median": 0.145056 + }, + "inter_kernel_gap_ms": { + "median": 0.040864 + }, + "kernel_sum_ms": { + "median": 0.008928 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 9859, + "submission_ms": { + "median": 0.145056 + }, + "synchronized_e2e_ms": { + "median": 0.165665 + } + }, + "baseline_07cf_precomputed_bench_iters": 14881, + "baseline_07cf_precomputed_gpu_span_ms": 0.007712, + "baseline_07cf_precomputed_host_enqueue_ms": 0.041312, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.007712, + "baseline_07cf_precomputed_submission_ms": 0.041312, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.056, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.007712 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.045984, + "synchronized_e2e_ms": 0.061472 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.007712 + }, + "host_enqueue_ms": { + "median": 0.041312 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.007712 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 14881, + "submission_ms": { + "median": 0.041312 + }, + "synchronized_e2e_ms": { + "median": 0.056 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 6.45643153526971, + "submission": 3.5112316034082105, + "synchronized_e2e": 2.9583035714285715 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.91092, + "after_init_synchronized_e2e_ms_per_call": 7.943624, + "including_init_host_enqueue_ms_per_call": 42.138827, + "including_init_synchronized_e2e_ms_per_call": 42.289067, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9216424, + "after_init_synchronized_e2e_ms_per_call": 0.9434609, + "including_init_host_enqueue_ms_per_call": 4.3444331, + "including_init_synchronized_e2e_ms_per_call": 4.3780052000000005, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22271464, + "after_init_synchronized_e2e_ms_per_call": 0.24344459000000002, + "including_init_host_enqueue_ms_per_call": 0.56499371, + "including_init_synchronized_e2e_ms_per_call": 0.58689902, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.152821864, + "after_init_synchronized_e2e_ms_per_call": 0.173442959, + "including_init_host_enqueue_ms_per_call": 0.187049771, + "including_init_synchronized_e2e_ms_per_call": 0.207788402, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.91092, + "synchronized_e2e_ms": 7.943624, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 8.749545, + "median": 0.049792, + "min": 0.041248, + "p90": 0.072416, + "sample_count": 9859 + }, + "host_enqueue_ms": { + "max": 75.01403, + "median": 0.145056, + "min": 0.121792, + "p90": 0.2152000000000001, + "sample_count": 9859 + }, + "sample_count": 9859, + "synchronized_e2e_ms": { + "max": 75.163854, + "median": 0.165665, + "min": 0.139008, + "p90": 0.23987860000000003, + "sample_count": 9859 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 14881, + "candidate_precomputed_gpu_span_ms": 0.00672, + "candidate_precomputed_host_enqueue_ms": 0.037696, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.00672, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.037696, + "candidate_precomputed_synchronized_e2e_ms": 0.04976, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.00672 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.033824, + "synchronized_e2e_ms": 0.050432 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.00672 + }, + "host_enqueue_ms": { + "median": 0.037696 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.00672 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 14881, + "submission_ms": { + "median": 0.037696 + }, + "synchronized_e2e_ms": { + "median": 0.04976 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282c4530", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc0481a840" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.5047619047619045, + "submission": 1.1103565365025465, + "synchronized_e2e": 1.1299035369774921 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 9859, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.010112, + "candidate_public_raw_host_enqueue_ms": 0.041856, + "candidate_public_raw_inter_kernel_gap_ms": 0.000192, + "candidate_public_raw_kernel_sum_ms": 0.009888, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.041856, + "candidate_public_raw_synchronized_e2e_ms": 0.056224, + "candidate_public_raw_tflops_from_gpu_span": 1.451746835443038, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.009888 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.074816, + "synchronized_e2e_ms": 0.096096 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.010112 + }, + "host_enqueue_ms": { + "median": 0.041856 + }, + "inter_kernel_gap_ms": { + "median": 0.000192 + }, + "kernel_sum_ms": { + "median": 0.009888 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 9859, + "submission_ms": { + "median": 0.041856 + }, + "synchronized_e2e_ms": { + "median": 0.056224 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.583523, + "after_init_synchronized_e2e_ms_per_call": 2.611971, + "including_init_host_enqueue_ms_per_call": 37.161063, + "including_init_synchronized_e2e_ms_per_call": 450.69976299999996, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2960227, + "after_init_synchronized_e2e_ms_per_call": 0.31179870000000004, + "including_init_host_enqueue_ms_per_call": 3.7537766999999995, + "including_init_synchronized_e2e_ms_per_call": 45.120577899999994, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06727266999999999, + "after_init_synchronized_e2e_ms_per_call": 0.08178147000000001, + "including_init_host_enqueue_ms_per_call": 0.41304807, + "including_init_synchronized_e2e_ms_per_call": 4.562659389999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.044397667, + "after_init_synchronized_e2e_ms_per_call": 0.058779747, + "including_init_host_enqueue_ms_per_call": 0.07897520699999999, + "including_init_synchronized_e2e_ms_per_call": 0.5068675389999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.583523, + "synchronized_e2e_ms": 2.611971, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.011136, + "median": 0.010112, + "min": 0.009792, + "p90": 0.010368, + "sample_count": 9859 + }, + "host_enqueue_ms": { + "max": 76.683535, + "median": 0.041856, + "min": 0.032224, + "p90": 0.065184, + "sample_count": 9859 + }, + "sample_count": 9859, + "synchronized_e2e_ms": { + "max": 76.862703, + "median": 0.056224, + "min": 0.045664, + "p90": 0.08416000000000001, + "sample_count": 9859 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.071072, + "submission_ms": 0.071072, + "synchronized_e2e_ms": 0.091584 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.045984, + "submission_ms": 0.045984, + "synchronized_e2e_ms": 0.061472 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.210017, + "submission_ms": 0.210017, + "synchronized_e2e_ms": 0.257057 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.91092, + "submission_ms": 7.91092, + "synchronized_e2e_ms": 7.943624 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.061024, + "submission_ms": 0.061024, + "synchronized_e2e_ms": 0.081184 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.033824, + "submission_ms": 0.033824, + "synchronized_e2e_ms": 0.050432 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.182049, + "submission_ms": 1.182049, + "synchronized_e2e_ms": 1.211361 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.055713, + "submission_ms": 1.055713, + "synchronized_e2e_ms": 1.083425 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.074816, + "submission_ms": 0.074816, + "synchronized_e2e_ms": 0.096096 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.583523, + "submission_ms": 2.583523, + "synchronized_e2e_ms": 2.611971 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.416031, + "evolution_kernel_ms": 0.166575, + "evolution_speedup": 2.4976, + "evolution_tflops": 0.0881, + "expected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d112_b1_n256_k256_d112", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 49480, + "measurement_schedule_sha256": "5b699916532ded30ec5e25a2333fefe51c8c724edf19c0dd657fb83888a4ab31", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 14881, + "public_pair_count": 9859, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 14881, + "baseline_public_raw": 9859, + "candidate_precomputed": 14881, + "candidate_public_raw": 9859 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 9898 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:post_d895_d112_b1_n256_k256_d112", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.1476190476190475, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.9465175014228797, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.041237440997622, + "including_init_synchronized_e2e_speedup": 0.09382979640040327, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 3.025865406109775, + "including_init_synchronized_e2e_speedup": 0.09702901433804555, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.976769554276782, + "including_init_synchronized_e2e_speedup": 0.12863090794949744, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.950726531708277, + "including_init_synchronized_e2e_speedup": 0.40994616149605123, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 4.924050632911393, + "hot_synchronized_e2e_speedup": 2.9465175014228797, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 211204, + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d112_b1_n256_k256_d112", + "source": "lowmid_small_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 112, + "K": 8192, + "N": 512, + "baseline_07cf_adapter_bench_iters": 1620, + "baseline_07cf_adapter_gpu_span_ms": 0.133088, + "baseline_07cf_adapter_host_enqueue_ms": 0.14590399999999998, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.037216, + "baseline_07cf_adapter_kernel_sum_ms": 0.095872, + "baseline_07cf_adapter_submission_ms": 0.14590399999999998, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.2333605, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.095872 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.222272, + "synchronized_e2e_ms": 0.306752 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.133088 + }, + "host_enqueue_ms": { + "median": 0.14590399999999998 + }, + "inter_kernel_gap_ms": { + "median": 0.037216 + }, + "kernel_sum_ms": { + "median": 0.095872 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1620, + "submission_ms": { + "median": 0.14590399999999998 + }, + "synchronized_e2e_ms": { + "median": 0.2333605 + } + }, + "baseline_07cf_precomputed_bench_iters": 1842, + "baseline_07cf_precomputed_gpu_span_ms": 0.101216, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042752, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.101216, + "baseline_07cf_precomputed_submission_ms": 0.042752, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.151472, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.101216 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.0456, + "synchronized_e2e_ms": 0.130112 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.101216 + }, + "host_enqueue_ms": { + "median": 0.042752 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.101216 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1842, + "submission_ms": { + "median": 0.042752 + }, + "synchronized_e2e_ms": { + "median": 0.151472 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.3148909263357573, + "submission": 3.4127994011976046, + "synchronized_e2e": 1.5406180680257737 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.554921, + "after_init_synchronized_e2e_ms_per_call": 8.615657, + "including_init_host_enqueue_ms_per_call": 43.004365, + "including_init_synchronized_e2e_ms_per_call": 492.657151, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9868057, + "after_init_synchronized_e2e_ms_per_call": 1.07159015, + "including_init_host_enqueue_ms_per_call": 4.4317501, + "including_init_synchronized_e2e_ms_per_call": 49.47573955, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22999416999999997, + "after_init_synchronized_e2e_ms_per_call": 0.31718346500000005, + "including_init_host_enqueue_ms_per_call": 0.57448861, + "including_init_synchronized_e2e_ms_per_call": 5.157598405, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15431301699999997, + "after_init_synchronized_e2e_ms_per_call": 0.2417427965, + "including_init_host_enqueue_ms_per_call": 0.18876246099999996, + "including_init_synchronized_e2e_ms_per_call": 0.7257842905, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.554921, + "synchronized_e2e_ms": 8.615657, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.571904, + "median": 0.133088, + "min": 0.124288, + "p90": 0.1401961, + "sample_count": 1620 + }, + "host_enqueue_ms": { + "max": 0.618912, + "median": 0.14590399999999998, + "min": 0.124673, + "p90": 0.16999039999999999, + "sample_count": 1620 + }, + "sample_count": 1620, + "synchronized_e2e_ms": { + "max": 0.685984, + "median": 0.2333605, + "min": 0.213633, + "p90": 0.2567041, + "sample_count": 1620 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 1842, + "candidate_precomputed_gpu_span_ms": 0.054016, + "candidate_precomputed_host_enqueue_ms": 0.038944, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.054016, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.038944, + "candidate_precomputed_synchronized_e2e_ms": 0.090464, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.054016 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.032384, + "synchronized_e2e_ms": 0.08592 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.054016 + }, + "host_enqueue_ms": { + "median": 0.038944 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.054016 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1842, + "submission_ms": { + "median": 0.038944 + }, + "synchronized_e2e_ms": { + "median": 0.090464 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc0450b500", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04508110" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.143957345971564, + "submission": 1.0706655710764175, + "synchronized_e2e": 1.166607711354793 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 1620, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.061792, + "candidate_public_raw_host_enqueue_ms": 0.041696, + "candidate_public_raw_inter_kernel_gap_ms": 0.00016, + "candidate_public_raw_kernel_sum_ms": 0.061601, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.041696, + "candidate_public_raw_synchronized_e2e_ms": 0.105536, + "candidate_public_raw_tflops_from_gpu_span": 15.204623511134127, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.061601 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.05744, + "synchronized_e2e_ms": 0.11248 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.061792 + }, + "host_enqueue_ms": { + "median": 0.041696 + }, + "inter_kernel_gap_ms": { + "median": 0.00016 + }, + "kernel_sum_ms": { + "median": 0.061601 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1620, + "submission_ms": { + "median": 0.041696 + }, + "synchronized_e2e_ms": { + "median": 0.105536 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 42.436236, + "after_init_synchronized_e2e_ms_per_call": 42.470636, + "including_init_host_enqueue_ms_per_call": 77.31031999999999, + "including_init_synchronized_e2e_ms_per_call": 77.413712, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 4.28115, + "after_init_synchronized_e2e_ms_per_call": 4.342046, + "including_init_host_enqueue_ms_per_call": 7.768558399999999, + "including_init_synchronized_e2e_ms_per_call": 7.836353600000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.46564140000000004, + "after_init_synchronized_e2e_ms_per_call": 0.529187, + "including_init_host_enqueue_ms_per_call": 0.8143822399999999, + "including_init_synchronized_e2e_ms_per_call": 0.87861776, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.08409054, + "after_init_synchronized_e2e_ms_per_call": 0.14790109999999998, + "including_init_host_enqueue_ms_per_call": 0.11896462399999999, + "including_init_synchronized_e2e_ms_per_call": 0.182844176, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 42.436236, + "synchronized_e2e_ms": 42.470636, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.062784, + "median": 0.061792, + "min": 0.061216, + "p90": 0.06208, + "sample_count": 1620 + }, + "host_enqueue_ms": { + "max": 0.075744, + "median": 0.041696, + "min": 0.034144, + "p90": 0.0496, + "sample_count": 1620 + }, + "sample_count": 1620, + "synchronized_e2e_ms": { + "max": 0.138304, + "median": 0.105536, + "min": 0.099456, + "p90": 0.11251520000000001, + "sample_count": 1620 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.064416, + "submission_ms": 0.064416, + "synchronized_e2e_ms": 0.14784 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.0456, + "submission_ms": 0.0456, + "synchronized_e2e_ms": 0.130112 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.222272, + "submission_ms": 0.222272, + "synchronized_e2e_ms": 0.306752 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.554921, + "submission_ms": 8.554921, + "synchronized_e2e_ms": 8.615657 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.062528, + "submission_ms": 0.062528, + "synchronized_e2e_ms": 0.10448 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.032384, + "submission_ms": 0.032384, + "synchronized_e2e_ms": 0.08592 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.122498, + "submission_ms": 1.122498, + "synchronized_e2e_ms": 1.14685 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.006881, + "submission_ms": 1.006881, + "synchronized_e2e_ms": 1.028737 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.05744, + "submission_ms": 0.05744, + "synchronized_e2e_ms": 0.11248 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 42.436236, + "submission_ms": 42.436236, + "synchronized_e2e_ms": 42.470636 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 1.731693, + "evolution_kernel_ms": 0.294304, + "evolution_speedup": 5.884, + "evolution_tflops": 3.1924, + "expected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d112_b1_n512_k8192_d112", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 6924, + "measurement_schedule_sha256": "ff82bc506ad0b3fdc9975dc1bb5f94cad411ace0df5324cd3e62a90709fa133a", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 1842, + "public_pair_count": 1620, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 1842, + "baseline_public_raw": 1620, + "candidate_precomputed": 1842, + "candidate_public_raw": 1620 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 1386 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:post_d895_d112_b1_n512_k8192_d112", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.8738151658767772, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.2111933368708305, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.20286150176795095, + "including_init_synchronized_e2e_speedup": 6.363952047668247, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.24679382714968937, + "including_init_synchronized_e2e_speedup": 6.313617541454484, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.5993787923739625, + "including_init_synchronized_e2e_speedup": 5.870127648000195, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.6344895102200052, + "including_init_synchronized_e2e_speedup": 3.9694143197648257, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.153806317969964, + "hot_synchronized_e2e_speedup": 2.2111933368708305, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 211203, + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d112_b1_n512_k8192_d112", + "source": "high_k_low_n", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 128, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 4, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 112, + "K": 512, + "N": 2048, + "baseline_07cf_adapter_bench_iters": 5140, + "baseline_07cf_adapter_gpu_span_ms": 0.064928, + "baseline_07cf_adapter_host_enqueue_ms": 0.171984, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.0486085, + "baseline_07cf_adapter_kernel_sum_ms": 0.01632, + "baseline_07cf_adapter_submission_ms": 0.171984, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.1938245, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.01632 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.256673, + "synchronized_e2e_ms": 0.279617 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.064928 + }, + "host_enqueue_ms": { + "median": 0.171984 + }, + "inter_kernel_gap_ms": { + "median": 0.0486085 + }, + "kernel_sum_ms": { + "median": 0.01632 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 5140, + "submission_ms": { + "median": 0.171984 + }, + "synchronized_e2e_ms": { + "median": 0.1938245 + } + }, + "baseline_07cf_precomputed_bench_iters": 7912, + "baseline_07cf_precomputed_gpu_span_ms": 0.0128, + "baseline_07cf_precomputed_host_enqueue_ms": 0.051968, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.0128, + "baseline_07cf_precomputed_submission_ms": 0.051968, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.07136, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.0128 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.052544, + "synchronized_e2e_ms": 0.069952 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.0128 + }, + "host_enqueue_ms": { + "median": 0.051968 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.0128 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 7912, + "submission_ms": { + "median": 0.051968 + }, + "synchronized_e2e_ms": { + "median": 0.07136 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 5.0725, + "submission": 3.30942118226601, + "synchronized_e2e": 2.716150504484305 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.568809, + "after_init_synchronized_e2e_ms_per_call": 8.596713, + "including_init_host_enqueue_ms_per_call": 44.120077, + "including_init_synchronized_e2e_ms_per_call": 44.22795, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 1.0116665, + "after_init_synchronized_e2e_ms_per_call": 1.03411335, + "including_init_host_enqueue_ms_per_call": 4.5667933000000005, + "including_init_synchronized_e2e_ms_per_call": 4.5972370499999995, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.25595225, + "after_init_synchronized_e2e_ms_per_call": 0.277853385, + "including_init_host_enqueue_ms_per_call": 0.61146493, + "including_init_synchronized_e2e_ms_per_call": 0.634165755, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.180380825, + "after_init_synchronized_e2e_ms_per_call": 0.20222738850000002, + "including_init_host_enqueue_ms_per_call": 0.21593209300000002, + "including_init_synchronized_e2e_ms_per_call": 0.2378586255, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.568809, + "synchronized_e2e_ms": 8.596713, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 70.985259, + "median": 0.064928, + "min": 0.052992, + "p90": 0.08906560000000002, + "sample_count": 5140 + }, + "host_enqueue_ms": { + "max": 86.131289, + "median": 0.171984, + "min": 0.144416, + "p90": 0.2166113, + "sample_count": 5140 + }, + "sample_count": 5140, + "synchronized_e2e_ms": { + "max": 86.334234, + "median": 0.1938245, + "min": 0.16368, + "p90": 0.24266330000000003, + "sample_count": 5140 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 7912, + "candidate_precomputed_gpu_span_ms": 0.014144, + "candidate_precomputed_host_enqueue_ms": 0.04688, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.014144, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.04688, + "candidate_precomputed_synchronized_e2e_ms": 0.059648, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.014144 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.048, + "synchronized_e2e_ms": 0.06432 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.014144 + }, + "host_enqueue_ms": { + "median": 0.04688 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.014144 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 7912, + "submission_ms": { + "median": 0.04688 + }, + "synchronized_e2e_ms": { + "median": 0.059648 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc009006e0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc00903860" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.3755656108597285, + "submission": 1.2245733788395905, + "synchronized_e2e": 1.3100858369098713 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 5140, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.019456, + "candidate_public_raw_host_enqueue_ms": 0.057408, + "candidate_public_raw_inter_kernel_gap_ms": 0.00016, + "candidate_public_raw_kernel_sum_ms": 0.019296, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.057408, + "candidate_public_raw_synchronized_e2e_ms": 0.078144, + "candidate_public_raw_tflops_from_gpu_span": 24.144842105263155, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.019296 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.0856, + "synchronized_e2e_ms": 0.10592 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.019456 + }, + "host_enqueue_ms": { + "median": 0.057408 + }, + "inter_kernel_gap_ms": { + "median": 0.00016 + }, + "kernel_sum_ms": { + "median": 0.019296 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 5140, + "submission_ms": { + "median": 0.057408 + }, + "synchronized_e2e_ms": { + "median": 0.078144 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 43.016524, + "after_init_synchronized_e2e_ms_per_call": 43.05006, + "including_init_host_enqueue_ms_per_call": 78.854544, + "including_init_synchronized_e2e_ms_per_call": 527.150273, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 4.3533196, + "after_init_synchronized_e2e_ms_per_call": 4.3753356000000005, + "including_init_host_enqueue_ms_per_call": 7.9371216, + "including_init_synchronized_e2e_ms_per_call": 52.7853569, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.48699915999999993, + "after_init_synchronized_e2e_ms_per_call": 0.50786316, + "including_init_host_enqueue_ms_per_call": 0.84537936, + "including_init_synchronized_e2e_ms_per_call": 5.34886529, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.10036711599999999, + "after_init_synchronized_e2e_ms_per_call": 0.12111591600000002, + "including_init_host_enqueue_ms_per_call": 0.136205136, + "including_init_synchronized_e2e_ms_per_call": 0.605216129, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 43.016524, + "synchronized_e2e_ms": 43.05006, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.020544, + "median": 0.019456, + "min": 0.01888, + "p90": 0.019744, + "sample_count": 5140 + }, + "host_enqueue_ms": { + "max": 0.596161, + "median": 0.057408, + "min": 0.042048, + "p90": 0.07709440000000002, + "sample_count": 5140 + }, + "sample_count": 5140, + "synchronized_e2e_ms": { + "max": 0.763745, + "median": 0.078144, + "min": 0.063808, + "p90": 0.0981481, + "sample_count": 5140 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.078272, + "submission_ms": 0.078272, + "synchronized_e2e_ms": 0.097824 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.052544, + "submission_ms": 0.052544, + "synchronized_e2e_ms": 0.069952 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.256673, + "submission_ms": 0.256673, + "synchronized_e2e_ms": 0.279617 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.568809, + "submission_ms": 8.568809, + "synchronized_e2e_ms": 8.596713 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.107616, + "submission_ms": 0.107616, + "synchronized_e2e_ms": 0.129024 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.048, + "submission_ms": 0.048, + "synchronized_e2e_ms": 0.06432 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.042465, + "submission_ms": 1.042465, + "synchronized_e2e_ms": 1.091969 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 0.870081, + "submission_ms": 0.870081, + "synchronized_e2e_ms": 0.894465 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.0856, + "submission_ms": 0.0856, + "synchronized_e2e_ms": 0.10592 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 43.016524, + "submission_ms": 43.016524, + "synchronized_e2e_ms": 43.05006 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.312623, + "evolution_kernel_ms": 0.188, + "evolution_speedup": 1.6629, + "evolution_tflops": 2.4987, + "expected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d112_b2_n2048_k512_d112", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 26104, + "measurement_schedule_sha256": "c1b581be8e8b641f63e4c92b2f2aa68ea735ff72981e2fb7a5cffcd058313985", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 7912, + "public_pair_count": 5140, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 7912, + "baseline_public_raw": 5140, + "candidate_precomputed": 7912, + "candidate_public_raw": 5140 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 5222 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:post_d895_d112_b2_n2048_k512_d112", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.9049773755656109, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.480350378787879, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.19969108056992252, + "including_init_synchronized_e2e_speedup": 0.08390007985445927, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.23635063559467295, + "including_init_synchronized_e2e_speedup": 0.08709303716008407, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.5471028554227088, + "including_init_synchronized_e2e_speedup": 0.11856080133212703, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.6697011852678387, + "including_init_synchronized_e2e_speedup": 0.39301435322454864, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 3.3371710526315788, + "hot_synchronized_e2e_speedup": 2.480350378787879, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 211201, + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d112_b2_n2048_k512_d112", + "source": "new_lowmid_gap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 112, + "K": 1024, + "N": 8192, + "baseline_07cf_adapter_bench_iters": 2239, + "baseline_07cf_adapter_gpu_span_ms": 0.07888, + "baseline_07cf_adapter_host_enqueue_ms": 0.162944, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.02592, + "baseline_07cf_adapter_kernel_sum_ms": 0.052864, + "baseline_07cf_adapter_submission_ms": 0.162944, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.188672, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.052864 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.234017, + "synchronized_e2e_ms": 0.257377 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.07888 + }, + "host_enqueue_ms": { + "median": 0.162944 + }, + "inter_kernel_gap_ms": { + "median": 0.02592 + }, + "kernel_sum_ms": { + "median": 0.052864 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2239, + "submission_ms": { + "median": 0.162944 + }, + "synchronized_e2e_ms": { + "median": 0.188672 + } + }, + "baseline_07cf_precomputed_bench_iters": 3412, + "baseline_07cf_precomputed_gpu_span_ms": 0.030976, + "baseline_07cf_precomputed_host_enqueue_ms": 0.046976, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.030976, + "baseline_07cf_precomputed_submission_ms": 0.046976, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.08352, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.030976 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.07104, + "synchronized_e2e_ms": 0.104992 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.030976 + }, + "host_enqueue_ms": { + "median": 0.046976 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.030976 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3412, + "submission_ms": { + "median": 0.046976 + }, + "synchronized_e2e_ms": { + "median": 0.08352 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.5464876033057853, + "submission": 3.46866485013624, + "synchronized_e2e": 2.259003831417625 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.79796, + "after_init_synchronized_e2e_ms_per_call": 7.82356, + "including_init_host_enqueue_ms_per_call": 44.732398, + "including_init_synchronized_e2e_ms_per_call": 462.934879, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9264455999999999, + "after_init_synchronized_e2e_ms_per_call": 0.9521608, + "including_init_host_enqueue_ms_per_call": 4.6198894, + "including_init_synchronized_e2e_ms_per_call": 46.4632927, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23929416, + "after_init_synchronized_e2e_ms_per_call": 0.26502088, + "including_init_host_enqueue_ms_per_call": 0.6086385400000001, + "including_init_synchronized_e2e_ms_per_call": 4.81613407, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.170579016, + "after_init_synchronized_e2e_ms_per_call": 0.19630688799999998, + "including_init_host_enqueue_ms_per_call": 0.207513454, + "including_init_synchronized_e2e_ms_per_call": 0.6514182070000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.79796, + "synchronized_e2e_ms": 7.82356, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.440165, + "median": 0.07888, + "min": 0.068672, + "p90": 0.10757120000000002, + "sample_count": 2239 + }, + "host_enqueue_ms": { + "max": 16.144145, + "median": 0.162944, + "min": 0.133473, + "p90": 0.2579264000000001, + "sample_count": 2239 + }, + "sample_count": 2239, + "synchronized_e2e_ms": { + "max": 21.332534, + "median": 0.188672, + "min": 0.161024, + "p90": 0.2866176000000001, + "sample_count": 2239 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3412, + "candidate_precomputed_gpu_span_ms": 0.026016, + "candidate_precomputed_host_enqueue_ms": 0.058032, + "candidate_precomputed_inter_kernel_gap_ms": 0.003328, + "candidate_precomputed_kernel_sum_ms": 0.022688, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.058032, + "candidate_precomputed_synchronized_e2e_ms": 0.071072, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.022688 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.043328, + "synchronized_e2e_ms": 0.05872 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.026016 + }, + "host_enqueue_ms": { + "median": 0.058032 + }, + "inter_kernel_gap_ms": { + "median": 0.003328 + }, + "kernel_sum_ms": { + "median": 0.022688 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3412, + "submission_ms": { + "median": 0.058032 + }, + "synchronized_e2e_ms": { + "median": 0.071072 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc028afd70", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc028afb30" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.7170971709717098, + "submission": 0.8447752963881996, + "synchronized_e2e": 1.329581269698334 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2239, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.044672, + "candidate_public_raw_host_enqueue_ms": 0.049024, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.044704, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.049024, + "candidate_public_raw_synchronized_e2e_ms": 0.094496, + "candidate_public_raw_tflops_from_gpu_span": 168.25288252148997, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.044608 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.075776, + "synchronized_e2e_ms": 0.099744 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.044672 + }, + "host_enqueue_ms": { + "median": 0.049024 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.044704 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2239, + "submission_ms": { + "median": 0.049024 + }, + "synchronized_e2e_ms": { + "median": 0.094496 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 52.046038, + "after_init_synchronized_e2e_ms_per_call": 52.080822, + "including_init_host_enqueue_ms_per_call": 89.404061, + "including_init_synchronized_e2e_ms_per_call": 89.516445, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 5.2487254, + "after_init_synchronized_e2e_ms_per_call": 5.2931286, + "including_init_host_enqueue_ms_per_call": 8.9845277, + "including_init_synchronized_e2e_ms_per_call": 9.0366909, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.56899414, + "after_init_synchronized_e2e_ms_per_call": 0.61435926, + "including_init_host_enqueue_ms_per_call": 0.94257437, + "including_init_synchronized_e2e_ms_per_call": 0.98871549, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.101021014, + "after_init_synchronized_e2e_ms_per_call": 0.146482326, + "including_init_host_enqueue_ms_per_call": 0.13837903699999998, + "including_init_synchronized_e2e_ms_per_call": 0.18391794900000002, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 52.046038, + "synchronized_e2e_ms": 52.080822, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.04592, + "median": 0.044672, + "min": 0.044128, + "p90": 0.044864, + "sample_count": 2239 + }, + "host_enqueue_ms": { + "max": 50.440501, + "median": 0.049024, + "min": 0.038304, + "p90": 0.08092160000000002, + "sample_count": 2239 + }, + "sample_count": 2239, + "synchronized_e2e_ms": { + "max": 50.611829, + "median": 0.094496, + "min": 0.083968, + "p90": 0.12261760000000002, + "sample_count": 2239 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.06912, + "submission_ms": 0.06912, + "synchronized_e2e_ms": 0.101856 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.07104, + "submission_ms": 0.07104, + "synchronized_e2e_ms": 0.104992 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.234017, + "submission_ms": 0.234017, + "synchronized_e2e_ms": 0.257377 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.79796, + "submission_ms": 7.79796, + "synchronized_e2e_ms": 7.82356 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.072929, + "submission_ms": 0.072929, + "synchronized_e2e_ms": 0.091745 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.043328, + "submission_ms": 0.043328, + "synchronized_e2e_ms": 0.05872 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.401345, + "submission_ms": 1.401345, + "synchronized_e2e_ms": 1.429505 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.208065, + "submission_ms": 1.208065, + "synchronized_e2e_ms": 1.232033 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.075776, + "submission_ms": 0.075776, + "synchronized_e2e_ms": 0.099744 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 52.046038, + "submission_ms": 52.046038, + "synchronized_e2e_ms": 52.080822 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.415103, + "evolution_kernel_ms": 0.1956, + "evolution_speedup": 2.1222, + "evolution_tflops": 38.4264, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d112_b4_n8192_k1024_d112", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 11302, + "measurement_schedule_sha256": "57ca68d2faf81cd6d1fe5433cd251d746c2cb07663c3f891543be5ab61b59aed", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3412, + "public_pair_count": 2239, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3412, + "baseline_public_raw": 2239, + "candidate_precomputed": 3412, + "candidate_public_raw": 2239 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2262 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:post_d895_d112_b4_n8192_k1024_d112", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.190651906519065, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.9966136132746362, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.15021959522835487, + "including_init_synchronized_e2e_speedup": 5.171506520394102, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.17988620189579374, + "including_init_synchronized_e2e_speedup": 5.141626864763073, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.4313776925898375, + "including_init_synchronized_e2e_speedup": 4.8711020700201635, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.3401404344166408, + "including_init_synchronized_e2e_speedup": 3.5418957776655065, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.765759312320917, + "hot_synchronized_e2e_speedup": 1.9966136132746362, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 211202, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d112_b4_n8192_k1024_d112", + "source": "new_lowmid_gap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 3, + "D": 128, + "K": 256, + "N": 1920, + "baseline_07cf_adapter_bench_iters": 8224, + "baseline_07cf_adapter_gpu_span_ms": 0.050848, + "baseline_07cf_adapter_host_enqueue_ms": 0.14761649999999998, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.038192500000000004, + "baseline_07cf_adapter_kernel_sum_ms": 0.01264, + "baseline_07cf_adapter_submission_ms": 0.14761649999999998, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.168192, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.01264 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.195328, + "synchronized_e2e_ms": 0.217856 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.050848 + }, + "host_enqueue_ms": { + "median": 0.14761649999999998 + }, + "inter_kernel_gap_ms": { + "median": 0.038192500000000004 + }, + "kernel_sum_ms": { + "median": 0.01264 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 8224, + "submission_ms": { + "median": 0.14761649999999998 + }, + "synchronized_e2e_ms": { + "median": 0.168192 + } + }, + "baseline_07cf_precomputed_bench_iters": 19410, + "baseline_07cf_precomputed_gpu_span_ms": 0.007968, + "baseline_07cf_precomputed_host_enqueue_ms": 0.041088, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.007968, + "baseline_07cf_precomputed_submission_ms": 0.041088, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.055904, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.007968 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.039776, + "synchronized_e2e_ms": 0.07072 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.007968 + }, + "host_enqueue_ms": { + "median": 0.041088 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.007968 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 19410, + "submission_ms": { + "median": 0.041088 + }, + "synchronized_e2e_ms": { + "median": 0.055904 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 6.381526104417671, + "submission": 3.5926912967289715, + "synchronized_e2e": 3.00858614768174 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.143913, + "after_init_synchronized_e2e_ms_per_call": 8.179337, + "including_init_host_enqueue_ms_per_call": 42.37182, + "including_init_synchronized_e2e_ms_per_call": 42.52478000000001, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.94724615, + "after_init_synchronized_e2e_ms_per_call": 0.9693065000000001, + "including_init_host_enqueue_ms_per_call": 4.37003685, + "including_init_synchronized_e2e_ms_per_call": 4.403850800000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22757946499999995, + "after_init_synchronized_e2e_ms_per_call": 0.24830345, + "including_init_host_enqueue_ms_per_call": 0.5698585349999999, + "including_init_synchronized_e2e_ms_per_call": 0.5917578800000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15561279649999998, + "after_init_synchronized_e2e_ms_per_call": 0.176203145, + "including_init_host_enqueue_ms_per_call": 0.18984070349999996, + "including_init_synchronized_e2e_ms_per_call": 0.210548588, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.143913, + "synchronized_e2e_ms": 8.179337, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 7.401992, + "median": 0.050848, + "min": 0.041824, + "p90": 0.07887040000000001, + "sample_count": 8224 + }, + "host_enqueue_ms": { + "max": 53.926456, + "median": 0.14761649999999998, + "min": 0.1232, + "p90": 0.21326079999999986, + "sample_count": 8224 + }, + "sample_count": 8224, + "synchronized_e2e_ms": { + "max": 54.042328, + "median": 0.168192, + "min": 0.141057, + "p90": 0.23863679999999998, + "sample_count": 8224 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 19410, + "candidate_precomputed_gpu_span_ms": 0.005792, + "candidate_precomputed_host_enqueue_ms": 0.039456, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.005792, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.039456, + "candidate_precomputed_synchronized_e2e_ms": 0.051264, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.005792 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.030592, + "synchronized_e2e_ms": 0.044032 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.005792 + }, + "host_enqueue_ms": { + "median": 0.039456 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.005792 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 19410, + "submission_ms": { + "median": 0.039456 + }, + "synchronized_e2e_ms": { + "median": 0.051264 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb9fb45a90", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb9fb44a70" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.165745856353591, + "submission": 1.1362530413625305, + "synchronized_e2e": 1.1679151061173534 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 8224, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.012544, + "candidate_public_raw_host_enqueue_ms": 0.044832, + "candidate_public_raw_inter_kernel_gap_ms": 0.000192, + "candidate_public_raw_kernel_sum_ms": 0.012352, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.044832, + "candidate_public_raw_synchronized_e2e_ms": 0.059872, + "candidate_public_raw_tflops_from_gpu_span": 30.093061224489794, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.012352 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.258689, + "synchronized_e2e_ms": 0.282529 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.012544 + }, + "host_enqueue_ms": { + "median": 0.044832 + }, + "inter_kernel_gap_ms": { + "median": 0.000192 + }, + "kernel_sum_ms": { + "median": 0.012352 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 8224, + "submission_ms": { + "median": 0.044832 + }, + "synchronized_e2e_ms": { + "median": 0.059872 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.335202, + "after_init_synchronized_e2e_ms_per_call": 2.363746, + "including_init_host_enqueue_ms_per_call": 36.912742, + "including_init_synchronized_e2e_ms_per_call": 450.45153799999997, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.27386900000000003, + "after_init_synchronized_e2e_ms_per_call": 0.29025939999999995, + "including_init_host_enqueue_ms_per_call": 3.7316230000000004, + "including_init_synchronized_e2e_ms_per_call": 45.09903859999999, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.0677357, + "after_init_synchronized_e2e_ms_per_call": 0.08291074, + "including_init_host_enqueue_ms_per_call": 0.41351109999999996, + "including_init_synchronized_e2e_ms_per_call": 4.563788659999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.04712237, + "after_init_synchronized_e2e_ms_per_call": 0.062175874, + "including_init_host_enqueue_ms_per_call": 0.08169990999999999, + "including_init_synchronized_e2e_ms_per_call": 0.510263666, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.335202, + "synchronized_e2e_ms": 2.363746, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.013312, + "median": 0.012544, + "min": 0.011808, + "p90": 0.012832, + "sample_count": 8224 + }, + "host_enqueue_ms": { + "max": 80.998868, + "median": 0.044832, + "min": 0.036576, + "p90": 0.06979549999999998, + "sample_count": 8224 + }, + "sample_count": 8224, + "synchronized_e2e_ms": { + "max": 81.184724, + "median": 0.059872, + "min": 0.051744, + "p90": 0.09090559999999998, + "sample_count": 8224 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.071456, + "submission_ms": 0.071456, + "synchronized_e2e_ms": 0.08928 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.039776, + "submission_ms": 0.039776, + "synchronized_e2e_ms": 0.07072 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.195328, + "submission_ms": 0.195328, + "synchronized_e2e_ms": 0.217856 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.143913, + "submission_ms": 8.143913, + "synchronized_e2e_ms": 8.179337 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.055776, + "submission_ms": 0.055776, + "synchronized_e2e_ms": 0.073184 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.030592, + "submission_ms": 0.030592, + "synchronized_e2e_ms": 0.044032 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.249217, + "submission_ms": 1.249217, + "synchronized_e2e_ms": 1.271009 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.089537, + "submission_ms": 1.089537, + "synchronized_e2e_ms": 1.112961 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.258689, + "submission_ms": 0.258689, + "synchronized_e2e_ms": 0.282529 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.335202, + "submission_ms": 2.335202, + "synchronized_e2e_ms": 2.363746 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.281664, + "evolution_kernel_ms": 0.1436, + "evolution_speedup": 1.9614, + "evolution_tflops": 2.6287, + "expected_route": "aligned_weave_v10_fallback", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d128_fallback_b3_n1920_k256_d128", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 55268, + "measurement_schedule_sha256": "bfcb63da948d8c302436d0ca4132e13dbbf3f215479c20e6f4ba6c51ce759e76", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 19410, + "public_pair_count": 8224, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 19410, + "baseline_public_raw": 8224, + "candidate_precomputed": 19410, + "candidate_public_raw": 8224 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 11054 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:post_d895_d128_fallback_b3_n1920_k256_d128", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.3756906077348066, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.809192944949225, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.46032822477542, + "including_init_synchronized_e2e_speedup": 0.09440478367286652, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 3.3394491272289555, + "including_init_synchronized_e2e_speedup": 0.097648440780731, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.9948285348798, + "including_init_synchronized_e2e_speedup": 0.1296637342536366, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.833947215603274, + "including_init_synchronized_e2e_speedup": 0.4126270436821579, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 4.053571428571429, + "hot_synchronized_e2e_speedup": 2.809192944949225, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 212803, + "selected_route": "aligned_weave_v10_fallback", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d128_fallback_b3_n1920_k256_d128", + "source": "near_floor_fallback_d128", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 5, + "D": 128, + "K": 512, + "N": 2176, + "baseline_07cf_adapter_bench_iters": 5750, + "baseline_07cf_adapter_gpu_span_ms": 0.059008, + "baseline_07cf_adapter_host_enqueue_ms": 0.15676849999999998, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.036736, + "baseline_07cf_adapter_kernel_sum_ms": 0.022241, + "baseline_07cf_adapter_submission_ms": 0.15676849999999998, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.176624, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.022241 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.228032, + "synchronized_e2e_ms": 0.248032 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.059008 + }, + "host_enqueue_ms": { + "median": 0.15676849999999998 + }, + "inter_kernel_gap_ms": { + "median": 0.036736 + }, + "kernel_sum_ms": { + "median": 0.022241 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 5750, + "submission_ms": { + "median": 0.15676849999999998 + }, + "synchronized_e2e_ms": { + "median": 0.176624 + } + }, + "baseline_07cf_precomputed_bench_iters": 13889, + "baseline_07cf_precomputed_gpu_span_ms": 0.013888, + "baseline_07cf_precomputed_host_enqueue_ms": 0.0432, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.013888, + "baseline_07cf_precomputed_submission_ms": 0.0432, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.064352, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.013888 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.050016, + "synchronized_e2e_ms": 0.069184 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.013888 + }, + "host_enqueue_ms": { + "median": 0.0432 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.013888 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 13889, + "submission_ms": { + "median": 0.0432 + }, + "synchronized_e2e_ms": { + "median": 0.064352 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 4.248847926267281, + "submission": 3.6289004629629624, + "synchronized_e2e": 2.744654400795624 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.112425, + "after_init_synchronized_e2e_ms_per_call": 8.137673, + "including_init_host_enqueue_ms_per_call": 42.561869, + "including_init_synchronized_e2e_ms_per_call": 492.179167, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9523341500000001, + "after_init_synchronized_e2e_ms_per_call": 0.9727288999999999, + "including_init_host_enqueue_ms_per_call": 4.39727855, + "including_init_synchronized_e2e_ms_per_call": 49.3768783, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23632506499999997, + "after_init_synchronized_e2e_ms_per_call": 0.25623449, + "including_init_host_enqueue_ms_per_call": 0.580819505, + "including_init_synchronized_e2e_ms_per_call": 5.09664943, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.1647241565, + "after_init_synchronized_e2e_ms_per_call": 0.184585049, + "including_init_host_enqueue_ms_per_call": 0.1991736005, + "including_init_synchronized_e2e_ms_per_call": 0.668626543, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.112425, + "synchronized_e2e_ms": 8.137673, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.653921, + "median": 0.059008, + "min": 0.043232, + "p90": 0.08529280000000004, + "sample_count": 5750 + }, + "host_enqueue_ms": { + "max": 143.258933, + "median": 0.15676849999999998, + "min": 0.126304, + "p90": 0.24692800000000006, + "sample_count": 5750 + }, + "sample_count": 5750, + "synchronized_e2e_ms": { + "max": 153.609503, + "median": 0.176624, + "min": 0.143808, + "p90": 0.27264730000000004, + "sample_count": 5750 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 13889, + "candidate_precomputed_gpu_span_ms": 0.007488, + "candidate_precomputed_host_enqueue_ms": 0.042208, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.007488, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.042208, + "candidate_precomputed_synchronized_e2e_ms": 0.054432, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.007488 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.038432, + "synchronized_e2e_ms": 0.05568 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.007488 + }, + "host_enqueue_ms": { + "median": 0.042208 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.007488 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 13889, + "submission_ms": { + "median": 0.042208 + }, + "synchronized_e2e_ms": { + "median": 0.054432 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7089a3f0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb70898c50" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.3504273504273505, + "submission": 1.133434420015163, + "synchronized_e2e": 1.2345679012345678 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 5750, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.0176, + "candidate_public_raw_host_enqueue_ms": 0.04784, + "candidate_public_raw_inter_kernel_gap_ms": 0.00016, + "candidate_public_raw_kernel_sum_ms": 0.017408, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.04784, + "candidate_public_raw_synchronized_e2e_ms": 0.0672, + "candidate_public_raw_tflops_from_gpu_span": 81.02632727272726, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.017408 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.296576, + "synchronized_e2e_ms": 0.320768 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.0176 + }, + "host_enqueue_ms": { + "median": 0.04784 + }, + "inter_kernel_gap_ms": { + "median": 0.00016 + }, + "kernel_sum_ms": { + "median": 0.017408 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 5750, + "submission_ms": { + "median": 0.04784 + }, + "synchronized_e2e_ms": { + "median": 0.0672 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.504131, + "after_init_synchronized_e2e_ms_per_call": 3.531843, + "including_init_host_enqueue_ms_per_call": 38.378215, + "including_init_synchronized_e2e_ms_per_call": 38.474919, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.3934691, + "after_init_synchronized_e2e_ms_per_call": 0.41366429999999993, + "including_init_host_enqueue_ms_per_call": 3.8808774999999995, + "including_init_synchronized_e2e_ms_per_call": 3.9079718999999997, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.08240291, + "after_init_synchronized_e2e_ms_per_call": 0.10184642999999999, + "including_init_host_enqueue_ms_per_call": 0.43114374999999994, + "including_init_synchronized_e2e_ms_per_call": 0.45127719, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.051296291, + "after_init_synchronized_e2e_ms_per_call": 0.07066464299999999, + "including_init_host_enqueue_ms_per_call": 0.08617037500000001, + "including_init_synchronized_e2e_ms_per_call": 0.10560771899999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.504131, + "synchronized_e2e_ms": 3.531843, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.018464, + "median": 0.0176, + "min": 0.017024, + "p90": 0.017824, + "sample_count": 5750 + }, + "host_enqueue_ms": { + "max": 111.522676, + "median": 0.04784, + "min": 0.037888, + "p90": 0.0793281, + "sample_count": 5750 + }, + "sample_count": 5750, + "synchronized_e2e_ms": { + "max": 134.967756, + "median": 0.0672, + "min": 0.058272, + "p90": 0.10086810000000002, + "sample_count": 5750 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.069728, + "submission_ms": 0.069728, + "synchronized_e2e_ms": 0.088864 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.050016, + "submission_ms": 0.050016, + "synchronized_e2e_ms": 0.069184 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.228032, + "submission_ms": 0.228032, + "synchronized_e2e_ms": 0.248032 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.112425, + "submission_ms": 8.112425, + "synchronized_e2e_ms": 8.137673 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.089856, + "submission_ms": 0.089856, + "synchronized_e2e_ms": 0.109152 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.038432, + "submission_ms": 0.038432, + "synchronized_e2e_ms": 0.05568 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.434402, + "submission_ms": 1.434402, + "synchronized_e2e_ms": 1.457442 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.225857, + "submission_ms": 1.225857, + "synchronized_e2e_ms": 1.248577 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.296576, + "submission_ms": 0.296576, + "synchronized_e2e_ms": 0.320768 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.504131, + "submission_ms": 3.504131, + "synchronized_e2e_ms": 3.531843 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.309232, + "evolution_kernel_ms": 0.14432, + "evolution_speedup": 2.1427, + "evolution_tflops": 9.8813, + "expected_route": "aligned_weave_v10_fallback", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d128_fallback_b5_n2176_k512_d128", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 39278, + "measurement_schedule_sha256": "f35a5fb92b280147517a82a3aaa62072a13cd834f1ef125093c0e22d2d5d4760", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 13889, + "public_pair_count": 5750, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 13889, + "baseline_public_raw": 5750, + "candidate_precomputed": 13889, + "candidate_public_raw": 5750 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 7856 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:post_d895_d128_fallback_b5_n2176_k512_d128", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.8547008547008546, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.6283333333333334, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.3040868464424946, + "including_init_synchronized_e2e_speedup": 12.792208009586712, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.3514934694630405, + "including_init_synchronized_e2e_speedup": 12.634911295037716, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.5158907386346288, + "including_init_synchronized_e2e_speedup": 11.293833464084459, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.6121273831384113, + "including_init_synchronized_e2e_speedup": 6.331227956926142, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 3.3527272727272726, + "hot_synchronized_e2e_speedup": 2.6283333333333334, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 212804, + "selected_route": "aligned_weave_v10_fallback", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d128_fallback_b5_n2176_k512_d128", + "source": "near_floor_fallback_d128", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 7, + "D": 128, + "K": 1024, + "N": 2432, + "baseline_07cf_adapter_bench_iters": 3681, + "baseline_07cf_adapter_gpu_span_ms": 0.07664, + "baseline_07cf_adapter_host_enqueue_ms": 0.175745, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.038976, + "baseline_07cf_adapter_kernel_sum_ms": 0.037632, + "baseline_07cf_adapter_submission_ms": 0.175745, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.198241, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.037632 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.212609, + "synchronized_e2e_ms": 0.232609 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.07664 + }, + "host_enqueue_ms": { + "median": 0.175745 + }, + "inter_kernel_gap_ms": { + "median": 0.038976 + }, + "kernel_sum_ms": { + "median": 0.037632 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3681, + "submission_ms": { + "median": 0.175745 + }, + "synchronized_e2e_ms": { + "median": 0.198241 + } + }, + "baseline_07cf_precomputed_bench_iters": 8942, + "baseline_07cf_precomputed_gpu_span_ms": 0.02544, + "baseline_07cf_precomputed_host_enqueue_ms": 0.051712, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.02544, + "baseline_07cf_precomputed_submission_ms": 0.051712, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.083776, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.02544 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.052064, + "synchronized_e2e_ms": 0.079712 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.02544 + }, + "host_enqueue_ms": { + "median": 0.051712 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.02544 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 8942, + "submission_ms": { + "median": 0.051712 + }, + "synchronized_e2e_ms": { + "median": 0.083776 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 3.0125786163522013, + "submission": 3.398534189356436, + "synchronized_e2e": 2.3663220970206265 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.366408, + "after_init_synchronized_e2e_ms_per_call": 7.393768, + "including_init_host_enqueue_ms_per_call": 42.917676, + "including_init_synchronized_e2e_ms_per_call": 43.025005, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8948113, + "after_init_synchronized_e2e_ms_per_call": 0.9177937, + "including_init_host_enqueue_ms_per_call": 4.4499381, + "including_init_synchronized_e2e_ms_per_call": 4.4809174, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.24765163, + "after_init_synchronized_e2e_ms_per_call": 0.27019627, + "including_init_host_enqueue_ms_per_call": 0.60316431, + "including_init_synchronized_e2e_ms_per_call": 0.62650864, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.182935663, + "after_init_synchronized_e2e_ms_per_call": 0.20543652699999998, + "including_init_host_enqueue_ms_per_call": 0.218486931, + "including_init_synchronized_e2e_ms_per_call": 0.241067764, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.366408, + "synchronized_e2e_ms": 7.393768, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.519552, + "median": 0.07664, + "min": 0.063712, + "p90": 0.101536, + "sample_count": 3681 + }, + "host_enqueue_ms": { + "max": 76.113775, + "median": 0.175745, + "min": 0.146016, + "p90": 0.221856, + "sample_count": 3681 + }, + "sample_count": 3681, + "synchronized_e2e_ms": { + "max": 76.299215, + "median": 0.198241, + "min": 0.167169, + "p90": 0.24768, + "sample_count": 3681 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 8942, + "candidate_precomputed_gpu_span_ms": 0.011488, + "candidate_precomputed_host_enqueue_ms": 0.048352, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.011488, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.048352, + "candidate_precomputed_synchronized_e2e_ms": 0.060928, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.011488 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.079104, + "synchronized_e2e_ms": 0.099552 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.011488 + }, + "host_enqueue_ms": { + "median": 0.048352 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.011488 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 8942, + "submission_ms": { + "median": 0.048352 + }, + "synchronized_e2e_ms": { + "median": 0.060928 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc02717a40", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04a2cb00" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.3732590529247912, + "submission": 1.270681667769689, + "synchronized_e2e": 1.4747899159663866 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 3681, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.027264, + "candidate_public_raw_host_enqueue_ms": 0.06144, + "candidate_public_raw_inter_kernel_gap_ms": 0.00016, + "candidate_public_raw_kernel_sum_ms": 0.027104, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.06144, + "candidate_public_raw_synchronized_e2e_ms": 0.089856, + "candidate_public_raw_tflops_from_gpu_span": 163.68615962441314, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.027104 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.31904, + "synchronized_e2e_ms": 0.345088 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.027264 + }, + "host_enqueue_ms": { + "median": 0.06144 + }, + "inter_kernel_gap_ms": { + "median": 0.00016 + }, + "kernel_sum_ms": { + "median": 0.027104 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3681, + "submission_ms": { + "median": 0.06144 + }, + "synchronized_e2e_ms": { + "median": 0.089856 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 40.647786, + "after_init_synchronized_e2e_ms_per_call": 40.678858, + "including_init_host_enqueue_ms_per_call": 76.485806, + "including_init_synchronized_e2e_ms_per_call": 524.779071, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 4.120074600000001, + "after_init_synchronized_e2e_ms_per_call": 4.148756199999999, + "including_init_host_enqueue_ms_per_call": 7.703876599999999, + "including_init_synchronized_e2e_ms_per_call": 52.558777500000005, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.46730346000000006, + "after_init_synchronized_e2e_ms_per_call": 0.49574602, + "including_init_host_enqueue_ms_per_call": 0.8256836599999999, + "including_init_synchronized_e2e_ms_per_call": 5.336748150000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.102026346, + "after_init_synchronized_e2e_ms_per_call": 0.130445002, + "including_init_host_enqueue_ms_per_call": 0.137864366, + "including_init_synchronized_e2e_ms_per_call": 0.6145452150000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 40.647786, + "synchronized_e2e_ms": 40.678858, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.027968, + "median": 0.027264, + "min": 0.02672, + "p90": 0.02752, + "sample_count": 3681 + }, + "host_enqueue_ms": { + "max": 81.462933, + "median": 0.06144, + "min": 0.049121, + "p90": 0.08304, + "sample_count": 3681 + }, + "sample_count": 3681, + "synchronized_e2e_ms": { + "max": 81.724021, + "median": 0.089856, + "min": 0.078528, + "p90": 0.109216, + "sample_count": 3681 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.080096, + "submission_ms": 0.080096, + "synchronized_e2e_ms": 0.106208 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.052064, + "submission_ms": 0.052064, + "synchronized_e2e_ms": 0.079712 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.212609, + "submission_ms": 0.212609, + "synchronized_e2e_ms": 0.232609 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.366408, + "submission_ms": 7.366408, + "synchronized_e2e_ms": 7.393768 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.084288, + "submission_ms": 0.084288, + "synchronized_e2e_ms": 0.102496 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.079104, + "submission_ms": 0.079104, + "synchronized_e2e_ms": 0.099552 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.246785, + "submission_ms": 1.246785, + "synchronized_e2e_ms": 1.269345 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.064705, + "submission_ms": 1.064705, + "synchronized_e2e_ms": 1.085633 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.31904, + "submission_ms": 0.31904, + "synchronized_e2e_ms": 0.345088 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 40.647786, + "submission_ms": 40.647786, + "synchronized_e2e_ms": 40.678858 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.4056, + "evolution_kernel_ms": 0.147936, + "evolution_speedup": 2.7417, + "evolution_tflops": 30.1667, + "expected_route": "aligned_weave_v10_fallback", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d128_fallback_b7_n2432_k1024_d128", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 25246, + "measurement_schedule_sha256": "bafdc295a607582a4f8a915ec52de926214008edc8dd53b98259047253e4171b", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 8942, + "public_pair_count": 3681, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 8942, + "baseline_public_raw": 3681, + "candidate_precomputed": 8942, + "candidate_public_raw": 3681 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 5052 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:post_d895_d128_fallback_b7_n2432_k1024_d128", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 2.214484679665738, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.20620771011396, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.18175947810530965, + "including_init_synchronized_e2e_speedup": 0.08198689196581926, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.22122141088936492, + "including_init_synchronized_e2e_speedup": 0.08525535815592361, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.5450296302933506, + "including_init_synchronized_e2e_speedup": 0.11739520441863081, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.574889983136341, + "including_init_synchronized_e2e_speedup": 0.3922701830816467, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.811032863849765, + "hot_synchronized_e2e_speedup": 2.20620771011396, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 212805, + "selected_route": "aligned_weave_v10_fallback", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d128_fallback_b7_n2432_k1024_d128", + "source": "near_floor_fallback_d128", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 128, + "K": 256, + "N": 262144, + "baseline_07cf_adapter_bench_iters": 268, + "baseline_07cf_adapter_gpu_span_ms": 0.377376, + "baseline_07cf_adapter_host_enqueue_ms": 0.218817, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.002272, + "baseline_07cf_adapter_kernel_sum_ms": 0.375104, + "baseline_07cf_adapter_submission_ms": 0.218817, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.5228809999999999, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.375104 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.352609, + "synchronized_e2e_ms": 0.582689 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.377376 + }, + "host_enqueue_ms": { + "median": 0.218817 + }, + "inter_kernel_gap_ms": { + "median": 0.002272 + }, + "kernel_sum_ms": { + "median": 0.375104 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 268, + "submission_ms": { + "median": 0.218817 + }, + "synchronized_e2e_ms": { + "median": 0.5228809999999999 + } + }, + "baseline_07cf_precomputed_bench_iters": 1190, + "baseline_07cf_precomputed_gpu_span_ms": 0.089152, + "baseline_07cf_precomputed_host_enqueue_ms": 0.057456, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.089152, + "baseline_07cf_precomputed_submission_ms": 0.057456, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.151408, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.089152 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.083712, + "synchronized_e2e_ms": 0.17168 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.089152 + }, + "host_enqueue_ms": { + "median": 0.057456 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.089152 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1190, + "submission_ms": { + "median": 0.057456 + }, + "synchronized_e2e_ms": { + "median": 0.151408 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 4.232950466618808, + "submission": 3.8084273182957395, + "synchronized_e2e": 3.453456884708866 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.998217, + "after_init_synchronized_e2e_ms_per_call": 9.052041, + "including_init_host_enqueue_ms_per_call": 45.932655, + "including_init_synchronized_e2e_ms_per_call": 464.16336, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 1.096757, + "after_init_synchronized_e2e_ms_per_call": 1.375797, + "including_init_host_enqueue_ms_per_call": 4.790200799999999, + "including_init_synchronized_e2e_ms_per_call": 46.8869289, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.306611, + "after_init_synchronized_e2e_ms_per_call": 0.6081726, + "including_init_host_enqueue_ms_per_call": 0.67595538, + "including_init_synchronized_e2e_ms_per_call": 5.15928579, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2275964, + "after_init_synchronized_e2e_ms_per_call": 0.5314101599999999, + "including_init_host_enqueue_ms_per_call": 0.264530838, + "including_init_synchronized_e2e_ms_per_call": 0.9865214789999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.998217, + "synchronized_e2e_ms": 9.052041, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.783713, + "median": 0.377376, + "min": 0.374912, + "p90": 0.37856, + "sample_count": 268 + }, + "host_enqueue_ms": { + "max": 0.809473, + "median": 0.218817, + "min": 0.140064, + "p90": 0.2736903, + "sample_count": 268 + }, + "sample_count": 268, + "synchronized_e2e_ms": { + "max": 0.878721, + "median": 0.5228809999999999, + "min": 0.473377, + "p90": 0.5626954, + "sample_count": 268 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 1190, + "candidate_precomputed_gpu_span_ms": 0.084256, + "candidate_precomputed_host_enqueue_ms": 0.0482405, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.084256, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.0482405, + "candidate_precomputed_synchronized_e2e_ms": 0.1240325, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.084256 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.053984, + "synchronized_e2e_ms": 0.125472 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.084256 + }, + "host_enqueue_ms": { + "median": 0.0482405 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.084256 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1190, + "submission_ms": { + "median": 0.0482405 + }, + "synchronized_e2e_ms": { + "median": 0.1240325 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb67415a30", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb674153a0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 4.435244967717432, + "submission": 1.4029705330583222, + "synchronized_e2e": 3.5300465603773206 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 268, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.373696, + "candidate_public_raw_host_enqueue_ms": 0.06767999999999999, + "candidate_public_raw_inter_kernel_gap_ms": 0.000192, + "candidate_public_raw_kernel_sum_ms": 0.373504, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.06767999999999999, + "candidate_public_raw_synchronized_e2e_ms": 0.4378405, + "candidate_public_raw_tflops_from_gpu_span": 91.94569481075528, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.373504 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.415968, + "synchronized_e2e_ms": 0.752961 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.373696 + }, + "host_enqueue_ms": { + "median": 0.06767999999999999 + }, + "inter_kernel_gap_ms": { + "median": 0.000192 + }, + "kernel_sum_ms": { + "median": 0.373504 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 268, + "submission_ms": { + "median": 0.06767999999999999 + }, + "synchronized_e2e_ms": { + "median": 0.4378405 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.031651, + "after_init_synchronized_e2e_ms_per_call": 3.365827, + "including_init_host_enqueue_ms_per_call": 40.389674, + "including_init_synchronized_e2e_ms_per_call": 40.80145, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.3640771, + "after_init_synchronized_e2e_ms_per_call": 0.73063915, + "including_init_host_enqueue_ms_per_call": 4.0998794, + "including_init_synchronized_e2e_ms_per_call": 4.474201450000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.09731970999999998, + "after_init_synchronized_e2e_ms_per_call": 0.46712036500000004, + "including_init_host_enqueue_ms_per_call": 0.47089993999999996, + "including_init_synchronized_e2e_ms_per_call": 0.841476595, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.070643971, + "after_init_synchronized_e2e_ms_per_call": 0.44076848650000006, + "including_init_host_enqueue_ms_per_call": 0.10800199399999999, + "including_init_synchronized_e2e_ms_per_call": 0.47820410950000003, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.031651, + "synchronized_e2e_ms": 3.365827, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.374977, + "median": 0.373696, + "min": 0.372128, + "p90": 0.374368, + "sample_count": 268 + }, + "host_enqueue_ms": { + "max": 0.152224, + "median": 0.06767999999999999, + "min": 0.0416, + "p90": 0.0877952, + "sample_count": 268 + }, + "sample_count": 268, + "synchronized_e2e_ms": { + "max": 0.508672, + "median": 0.4378405, + "min": 0.417952, + "p90": 0.453632, + "sample_count": 268 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.132577, + "submission_ms": 0.132577, + "synchronized_e2e_ms": 0.218177 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.083712, + "submission_ms": 0.083712, + "synchronized_e2e_ms": 0.17168 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.352609, + "submission_ms": 0.352609, + "synchronized_e2e_ms": 0.582689 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.998217, + "submission_ms": 8.998217, + "synchronized_e2e_ms": 9.052041 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.082464, + "submission_ms": 0.082464, + "synchronized_e2e_ms": 0.140064 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.053984, + "submission_ms": 0.053984, + "synchronized_e2e_ms": 0.125472 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.665409, + "submission_ms": 1.665409, + "synchronized_e2e_ms": 1.702626 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.458849, + "submission_ms": 1.458849, + "synchronized_e2e_ms": 1.492673 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.415968, + "submission_ms": 0.415968, + "synchronized_e2e_ms": 0.752961 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.031651, + "submission_ms": 3.031651, + "synchronized_e2e_ms": 3.365827 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.663294, + "evolution_kernel_ms": 0.24464, + "evolution_speedup": 2.7113, + "evolution_tflops": 140.4505, + "expected_route": "paired_large_v15", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d128_paired_b2_n262144_k256_d128", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 2916, + "measurement_schedule_sha256": "bc9ecdcaa6970b2ff0a467c44363b4c66b6db11400ab1440f4eab5ef118052e7", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 1190, + "public_pair_count": 268, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 1190, + "baseline_public_raw": 268, + "candidate_precomputed": 1190, + "candidate_public_raw": 268 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 584 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:post_d895_d128_paired_b2_n262144_k256_d128", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.058108621344474, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.1942271215202793, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.6893958007942773, + "including_init_synchronized_e2e_speedup": 11.376148641776211, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 1.8830047637058593, + "including_init_synchronized_e2e_speedup": 10.479396027194975, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 1.3019612193529604, + "including_init_synchronized_e2e_speedup": 6.131229104476756, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.2056446326727122, + "including_init_synchronized_e2e_speedup": 2.062971562564541, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.0098475766398356, + "hot_synchronized_e2e_speedup": 1.1942271215202793, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 212802, + "selected_route": "paired_large_v15", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d128_paired_b2_n262144_k256_d128", + "source": "near_floor_paired_d128", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 128, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 8, + "D": 128, + "K": 256, + "N": 8192, + "baseline_07cf_adapter_bench_iters": 1804, + "baseline_07cf_adapter_gpu_span_ms": 0.061472, + "baseline_07cf_adapter_host_enqueue_ms": 0.147584, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.004928, + "baseline_07cf_adapter_kernel_sum_ms": 0.056544, + "baseline_07cf_adapter_submission_ms": 0.147584, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.1680165, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.056544 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.189792, + "synchronized_e2e_ms": 0.20832 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.061472 + }, + "host_enqueue_ms": { + "median": 0.147584 + }, + "inter_kernel_gap_ms": { + "median": 0.004928 + }, + "kernel_sum_ms": { + "median": 0.056544 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1804, + "submission_ms": { + "median": 0.147584 + }, + "synchronized_e2e_ms": { + "median": 0.1680165 + } + }, + "baseline_07cf_precomputed_bench_iters": 5566, + "baseline_07cf_precomputed_gpu_span_ms": 0.018336, + "baseline_07cf_precomputed_host_enqueue_ms": 0.04032, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.018336, + "baseline_07cf_precomputed_submission_ms": 0.04032, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.065376, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.018336 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.041088, + "synchronized_e2e_ms": 0.064 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.018336 + }, + "host_enqueue_ms": { + "median": 0.04032 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.018336 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5566, + "submission_ms": { + "median": 0.04032 + }, + "synchronized_e2e_ms": { + "median": 0.065376 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 3.352530541012216, + "submission": 3.66031746031746, + "synchronized_e2e": 2.5700027533039647 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.333065, + "after_init_synchronized_e2e_ms_per_call": 8.359337, + "including_init_host_enqueue_ms_per_call": 42.560972, + "including_init_synchronized_e2e_ms_per_call": 42.70478, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9661320999999999, + "after_init_synchronized_e2e_ms_per_call": 0.98714855, + "including_init_host_enqueue_ms_per_call": 4.3889228000000005, + "including_init_synchronized_e2e_ms_per_call": 4.42169285, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22943880999999997, + "after_init_synchronized_e2e_ms_per_call": 0.24992970500000003, + "including_init_host_enqueue_ms_per_call": 0.57171788, + "including_init_synchronized_e2e_ms_per_call": 0.593384135, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.155769481, + "after_init_synchronized_e2e_ms_per_call": 0.17620782050000003, + "including_init_host_enqueue_ms_per_call": 0.18999738799999996, + "including_init_synchronized_e2e_ms_per_call": 0.2105532635, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.333065, + "synchronized_e2e_ms": 8.359337, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.225382, + "median": 0.061472, + "min": 0.057088, + "p90": 0.07968320000000001, + "sample_count": 1804 + }, + "host_enqueue_ms": { + "max": 61.305631, + "median": 0.147584, + "min": 0.12464, + "p90": 0.18805149999999998, + "sample_count": 1804 + }, + "sample_count": 1804, + "synchronized_e2e_ms": { + "max": 66.480612, + "median": 0.1680165, + "min": 0.145888, + "p90": 0.2130055, + "sample_count": 1804 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 5566, + "candidate_precomputed_gpu_span_ms": 0.017792, + "candidate_precomputed_host_enqueue_ms": 0.038592, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.017792, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.038592, + "candidate_precomputed_synchronized_e2e_ms": 0.051296, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.017792 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.033504, + "synchronized_e2e_ms": 0.047008 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.017792 + }, + "host_enqueue_ms": { + "median": 0.038592 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.017792 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 5566, + "submission_ms": { + "median": 0.038592 + }, + "synchronized_e2e_ms": { + "median": 0.051296 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04735af0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb9e8dade0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 3.129496402877698, + "submission": 1.1475953565505803, + "synchronized_e2e": 1.9806612601372426 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 1804, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.05568, + "candidate_public_raw_host_enqueue_ms": 0.044288, + "candidate_public_raw_inter_kernel_gap_ms": 0.000192, + "candidate_public_raw_kernel_sum_ms": 0.055488, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.044288, + "candidate_public_raw_synchronized_e2e_ms": 0.1016, + "candidate_public_raw_tflops_from_gpu_span": 77.13662528735632, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.055488 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.228416, + "synchronized_e2e_ms": 0.276576 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.05568 + }, + "host_enqueue_ms": { + "median": 0.044288 + }, + "inter_kernel_gap_ms": { + "median": 0.000192 + }, + "kernel_sum_ms": { + "median": 0.055488 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1804, + "submission_ms": { + "median": 0.044288 + }, + "synchronized_e2e_ms": { + "median": 0.1016 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.480003, + "after_init_synchronized_e2e_ms_per_call": 2.517827, + "including_init_host_enqueue_ms_per_call": 37.057542999999995, + "including_init_synchronized_e2e_ms_per_call": 450.605619, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2878595, + "after_init_synchronized_e2e_ms_per_call": 0.3432227, + "including_init_host_enqueue_ms_per_call": 3.7456134999999997, + "including_init_synchronized_e2e_ms_per_call": 45.1520019, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06864515, + "after_init_synchronized_e2e_ms_per_call": 0.12576226999999998, + "including_init_host_enqueue_ms_per_call": 0.41442054999999994, + "including_init_synchronized_e2e_ms_per_call": 4.60664019, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.046723715, + "after_init_synchronized_e2e_ms_per_call": 0.10401622699999999, + "including_init_host_enqueue_ms_per_call": 0.081301255, + "including_init_synchronized_e2e_ms_per_call": 0.552104019, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.480003, + "synchronized_e2e_ms": 2.517827, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.05664, + "median": 0.05568, + "min": 0.055072, + "p90": 0.056032, + "sample_count": 1804 + }, + "host_enqueue_ms": { + "max": 8.610665, + "median": 0.044288, + "min": 0.037408, + "p90": 0.062528, + "sample_count": 1804 + }, + "sample_count": 1804, + "synchronized_e2e_ms": { + "max": 8.654153, + "median": 0.1016, + "min": 0.095136, + "p90": 0.1177223, + "sample_count": 1804 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.063488, + "submission_ms": 0.063488, + "synchronized_e2e_ms": 0.085408 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.041088, + "submission_ms": 0.041088, + "synchronized_e2e_ms": 0.064 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.189792, + "submission_ms": 0.189792, + "synchronized_e2e_ms": 0.20832 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.333065, + "submission_ms": 8.333065, + "synchronized_e2e_ms": 8.359337 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.052576, + "submission_ms": 0.052576, + "synchronized_e2e_ms": 0.068864 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.033504, + "submission_ms": 0.033504, + "synchronized_e2e_ms": 0.047008 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.268514, + "submission_ms": 1.268514, + "synchronized_e2e_ms": 1.291938 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.077569, + "submission_ms": 1.077569, + "synchronized_e2e_ms": 1.098049 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.228416, + "submission_ms": 0.228416, + "synchronized_e2e_ms": 0.276576 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.480003, + "submission_ms": 2.480003, + "synchronized_e2e_ms": 2.517827 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.292671, + "evolution_kernel_ms": 0.15424, + "evolution_speedup": 1.8975, + "evolution_tflops": 27.846, + "expected_route": "d128_even_near_floor_v10_repair", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d128_paired_b8_n8192_k256_d128", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 14740, + "measurement_schedule_sha256": "726e9594c196113eaf4e32ca3774f9bb7943766eb79303af55c9839e7112e865", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 5566, + "public_pair_count": 1804, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 5566, + "baseline_public_raw": 1804, + "candidate_precomputed": 5566, + "candidate_public_raw": 1804 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2950 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:post_d895_d128_paired_b8_n8192_k256_d128", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.0305755395683456, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.6537057086614175, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.3200601153296074, + "including_init_synchronized_e2e_speedup": 0.09477196510503345, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.8761167312068814, + "including_init_synchronized_e2e_speedup": 0.09792905439260269, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 1.987318652883731, + "including_init_synchronized_e2e_speedup": 0.12881061045056352, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.6940416469826391, + "including_init_synchronized_e2e_speedup": 0.38136520701545557, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.104022988505747, + "hot_synchronized_e2e_speedup": 1.6537057086614175, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 212801, + "selected_route": "d128_even_near_floor_v10_repair", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d128_paired_b8_n8192_k256_d128", + "source": "near_floor_paired_d128", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 128, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 144, + "K": 256, + "N": 256, + "baseline_07cf_adapter_bench_iters": 9006, + "baseline_07cf_adapter_gpu_span_ms": 0.052736, + "baseline_07cf_adapter_host_enqueue_ms": 0.149856, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.042593, + "baseline_07cf_adapter_kernel_sum_ms": 0.010112, + "baseline_07cf_adapter_submission_ms": 0.149856, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.168864, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.010112 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.249184, + "synchronized_e2e_ms": 0.27008 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.052736 + }, + "host_enqueue_ms": { + "median": 0.149856 + }, + "inter_kernel_gap_ms": { + "median": 0.042593 + }, + "kernel_sum_ms": { + "median": 0.010112 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 9006, + "submission_ms": { + "median": 0.149856 + }, + "synchronized_e2e_ms": { + "median": 0.168864 + } + }, + "baseline_07cf_precomputed_bench_iters": 11141, + "baseline_07cf_precomputed_gpu_span_ms": 0.008992, + "baseline_07cf_precomputed_host_enqueue_ms": 0.04304, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.008992, + "baseline_07cf_precomputed_submission_ms": 0.04304, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.059456, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.008992 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.040032, + "synchronized_e2e_ms": 0.056032 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.008992 + }, + "host_enqueue_ms": { + "median": 0.04304 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.008992 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 11141, + "submission_ms": { + "median": 0.04304 + }, + "synchronized_e2e_ms": { + "median": 0.059456 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 5.864768683274021, + "submission": 3.4817843866171, + "synchronized_e2e": 2.840150699677072 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.983848, + "after_init_synchronized_e2e_ms_per_call": 8.008136, + "including_init_host_enqueue_ms_per_call": 42.433292, + "including_init_synchronized_e2e_ms_per_call": 492.04963, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9332552, + "after_init_synchronized_e2e_ms_per_call": 0.9527912000000001, + "including_init_host_enqueue_ms_per_call": 4.3781996, + "including_init_synchronized_e2e_ms_per_call": 49.356940599999994, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22819592, + "after_init_synchronized_e2e_ms_per_call": 0.24725671999999999, + "including_init_host_enqueue_ms_per_call": 0.57269036, + "including_init_synchronized_e2e_ms_per_call": 5.08767166, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.157689992, + "after_init_synchronized_e2e_ms_per_call": 0.176703272, + "including_init_host_enqueue_ms_per_call": 0.192139436, + "including_init_synchronized_e2e_ms_per_call": 0.660744766, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.983848, + "synchronized_e2e_ms": 8.008136, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 75.824048, + "median": 0.052736, + "min": 0.037344, + "p90": 0.076128, + "sample_count": 9006 + }, + "host_enqueue_ms": { + "max": 115.612184, + "median": 0.149856, + "min": 0.123552, + "p90": 0.2165925, + "sample_count": 9006 + }, + "sample_count": 9006, + "synchronized_e2e_ms": { + "max": 115.950552, + "median": 0.168864, + "min": 0.140224, + "p90": 0.24179250000000002, + "sample_count": 9006 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 11141, + "candidate_precomputed_gpu_span_ms": 0.015808, + "candidate_precomputed_host_enqueue_ms": 0.056704, + "candidate_precomputed_inter_kernel_gap_ms": 0.007072, + "candidate_precomputed_kernel_sum_ms": 0.008768, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.056704, + "candidate_precomputed_synchronized_e2e_ms": 0.069121, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.008768 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.0408, + "synchronized_e2e_ms": 0.055072 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.015808 + }, + "host_enqueue_ms": { + "median": 0.056704 + }, + "inter_kernel_gap_ms": { + "median": 0.007072 + }, + "kernel_sum_ms": { + "median": 0.008768 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 11141, + "submission_ms": { + "median": 0.056704 + }, + "synchronized_e2e_ms": { + "median": 0.069121 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04dc85f0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04dca510" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 0.7308324898785425, + "submission": 0.8205417607223476, + "synchronized_e2e": 0.8874871601973351 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 9006, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.011553, + "candidate_public_raw_host_enqueue_ms": 0.046528, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.01152, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.046528, + "candidate_public_raw_synchronized_e2e_ms": 0.061344, + "candidate_public_raw_tflops_from_gpu_span": 1.6337200727083874, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.01152 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.053216, + "synchronized_e2e_ms": 0.07152 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.011553 + }, + "host_enqueue_ms": { + "median": 0.046528 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.01152 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 9006, + "submission_ms": { + "median": 0.046528 + }, + "synchronized_e2e_ms": { + "median": 0.061344 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 70.445929, + "after_init_synchronized_e2e_ms_per_call": 70.475497, + "including_init_host_enqueue_ms_per_call": 105.320013, + "including_init_synchronized_e2e_ms_per_call": 105.41857300000001, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 7.0864681, + "after_init_synchronized_e2e_ms_per_call": 7.102759300000001, + "including_init_host_enqueue_ms_per_call": 10.5738765, + "including_init_synchronized_e2e_ms_per_call": 10.597066900000002, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.7505220100000001, + "after_init_synchronized_e2e_ms_per_call": 0.76548553, + "including_init_host_enqueue_ms_per_call": 1.0992628500000001, + "including_init_synchronized_e2e_ms_per_call": 1.11491629, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.116927401, + "after_init_synchronized_e2e_ms_per_call": 0.131758153, + "including_init_host_enqueue_ms_per_call": 0.151801485, + "including_init_synchronized_e2e_ms_per_call": 0.166701229, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 70.445929, + "synchronized_e2e_ms": 70.475497, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.01232, + "median": 0.011553, + "min": 0.010912, + "p90": 0.011968, + "sample_count": 9006 + }, + "host_enqueue_ms": { + "max": 124.737762, + "median": 0.046528, + "min": 0.036064, + "p90": 0.068608, + "sample_count": 9006 + }, + "sample_count": 9006, + "synchronized_e2e_ms": { + "max": 129.931175, + "median": 0.061344, + "min": 0.051616, + "p90": 0.08848, + "sample_count": 9006 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.064064, + "submission_ms": 0.064064, + "synchronized_e2e_ms": 0.081088 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.040032, + "submission_ms": 0.040032, + "synchronized_e2e_ms": 0.056032 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.249184, + "submission_ms": 0.249184, + "synchronized_e2e_ms": 0.27008 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.983848, + "submission_ms": 7.983848, + "synchronized_e2e_ms": 8.008136 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.081792, + "submission_ms": 0.081792, + "synchronized_e2e_ms": 0.09824 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.0408, + "submission_ms": 0.0408, + "synchronized_e2e_ms": 0.055072 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.336642, + "submission_ms": 1.336642, + "synchronized_e2e_ms": 1.36205 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.073889, + "submission_ms": 1.073889, + "synchronized_e2e_ms": 1.095201 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.053216, + "submission_ms": 0.053216, + "synchronized_e2e_ms": 0.07152 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 70.445929, + "submission_ms": 70.445929, + "synchronized_e2e_ms": 70.475497 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.537279, + "evolution_kernel_ms": 0.167728, + "evolution_speedup": 3.2033, + "evolution_tflops": 0.1125, + "expected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d144_b1_n256_k256_d144", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 40294, + "measurement_schedule_sha256": "946f147e2db31948590d15b00680dae30c4a2f86f935cc87c3e6d106f3bf11fc", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 11141, + "public_pair_count": 9006, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 11141, + "baseline_public_raw": 9006, + "candidate_precomputed": 11141, + "candidate_public_raw": 9006 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 8062 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:post_d895_d144_b1_n256_k256_d144", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.5688259109311741, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.7527386541471044, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.11363007486133797, + "including_init_synchronized_e2e_speedup": 4.667580066749717, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.1341438108426397, + "including_init_synchronized_e2e_speedup": 4.657603945106734, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.32300639307969675, + "including_init_synchronized_e2e_speedup": 4.563276817849705, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.3411183139460068, + "including_init_synchronized_e2e_speedup": 3.9636466387419373, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 4.564701809053925, + "hot_synchronized_e2e_speedup": 2.7527386541471044, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 214401, + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d144_b1_n256_k256_d144", + "source": "tailpad_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 144, + "K": 8192, + "N": 512, + "baseline_07cf_adapter_bench_iters": 889, + "baseline_07cf_adapter_gpu_span_ms": 0.170721, + "baseline_07cf_adapter_host_enqueue_ms": 0.175392, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.044, + "baseline_07cf_adapter_kernel_sum_ms": 0.126688, + "baseline_07cf_adapter_submission_ms": 0.175392, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.287584, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.126688 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.210304, + "synchronized_e2e_ms": 0.318016 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.170721 + }, + "host_enqueue_ms": { + "median": 0.175392 + }, + "inter_kernel_gap_ms": { + "median": 0.044 + }, + "kernel_sum_ms": { + "median": 0.126688 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 889, + "submission_ms": { + "median": 0.175392 + }, + "synchronized_e2e_ms": { + "median": 0.287584 + } + }, + "baseline_07cf_precomputed_bench_iters": 791, + "baseline_07cf_precomputed_gpu_span_ms": 0.129408, + "baseline_07cf_precomputed_host_enqueue_ms": 0.056288, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.129408, + "baseline_07cf_precomputed_submission_ms": 0.056288, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.192608, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.129408 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.05408, + "synchronized_e2e_ms": 0.168288 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.129408 + }, + "host_enqueue_ms": { + "median": 0.056288 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.129408 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 791, + "submission_ms": { + "median": 0.056288 + }, + "synchronized_e2e_ms": { + "median": 0.192608 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.3192461053412465, + "submission": 3.1159749857873793, + "synchronized_e2e": 1.4931051669712576 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.792904, + "after_init_synchronized_e2e_ms_per_call": 7.882312, + "including_init_host_enqueue_ms_per_call": 43.344172, + "including_init_synchronized_e2e_ms_per_call": 43.513549, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9371432000000001, + "after_init_synchronized_e2e_ms_per_call": 1.0470568, + "including_init_host_enqueue_ms_per_call": 4.4922699999999995, + "including_init_synchronized_e2e_ms_per_call": 4.6101805, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.25156712, + "after_init_synchronized_e2e_ms_per_call": 0.36353127999999996, + "including_init_host_enqueue_ms_per_call": 0.6070798, + "including_init_synchronized_e2e_ms_per_call": 0.71984365, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.18300951199999999, + "after_init_synchronized_e2e_ms_per_call": 0.29517872800000006, + "including_init_host_enqueue_ms_per_call": 0.21856077999999995, + "including_init_synchronized_e2e_ms_per_call": 0.33080996500000004, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.792904, + "synchronized_e2e_ms": 7.882312, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.449248, + "median": 0.170721, + "min": 0.158976, + "p90": 0.1850248, + "sample_count": 889 + }, + "host_enqueue_ms": { + "max": 0.786817, + "median": 0.175392, + "min": 0.14704, + "p90": 0.20280320000000002, + "sample_count": 889 + }, + "sample_count": 889, + "synchronized_e2e_ms": { + "max": 0.900993, + "median": 0.287584, + "min": 0.262656, + "p90": 0.3127424, + "sample_count": 889 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 791, + "candidate_precomputed_gpu_span_ms": 0.114528, + "candidate_precomputed_host_enqueue_ms": 0.066912, + "candidate_precomputed_inter_kernel_gap_ms": 0.007712, + "candidate_precomputed_kernel_sum_ms": 0.10672, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.066912, + "candidate_precomputed_synchronized_e2e_ms": 0.15888, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.10672 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.053888, + "synchronized_e2e_ms": 0.14368 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.114528 + }, + "host_enqueue_ms": { + "median": 0.066912 + }, + "inter_kernel_gap_ms": { + "median": 0.007712 + }, + "kernel_sum_ms": { + "median": 0.10672 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 791, + "submission_ms": { + "median": 0.066912 + }, + "synchronized_e2e_ms": { + "median": 0.15888 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc049f74d0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc049f6fc0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 0.9832355406538139, + "submission": 0.9320899091343855, + "synchronized_e2e": 1.1129909365558912 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 889, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.112608, + "candidate_public_raw_host_enqueue_ms": 0.062368, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.112832, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.062368, + "candidate_public_raw_synchronized_e2e_ms": 0.176832, + "candidate_public_raw_tflops_from_gpu_span": 10.72712020460358, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.112608 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.09328, + "synchronized_e2e_ms": 0.191808 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.112608 + }, + "host_enqueue_ms": { + "median": 0.062368 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.112832 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 889, + "submission_ms": { + "median": 0.062368 + }, + "synchronized_e2e_ms": { + "median": 0.176832 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 79.24757, + "after_init_synchronized_e2e_ms_per_call": 79.325138, + "including_init_host_enqueue_ms_per_call": 115.08559, + "including_init_synchronized_e2e_ms_per_call": 563.425351, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 7.9808882, + "after_init_synchronized_e2e_ms_per_call": 8.0916626, + "including_init_host_enqueue_ms_per_call": 11.5646902, + "including_init_synchronized_e2e_ms_per_call": 56.5016839, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.85422002, + "after_init_synchronized_e2e_ms_per_call": 0.9683150599999999, + "including_init_host_enqueue_ms_per_call": 1.21260022, + "including_init_synchronized_e2e_ms_per_call": 5.809317189999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.141553202, + "after_init_synchronized_e2e_ms_per_call": 0.255980306, + "including_init_host_enqueue_ms_per_call": 0.177391222, + "including_init_synchronized_e2e_ms_per_call": 0.7400805189999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 79.24757, + "synchronized_e2e_ms": 79.325138, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.113632, + "median": 0.112608, + "min": 0.111616, + "p90": 0.113024, + "sample_count": 889 + }, + "host_enqueue_ms": { + "max": 0.1008, + "median": 0.062368, + "min": 0.044352, + "p90": 0.0741696, + "sample_count": 889 + }, + "sample_count": 889, + "synchronized_e2e_ms": { + "max": 0.21168, + "median": 0.176832, + "min": 0.160192, + "p90": 0.186144, + "sample_count": 889 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.082528, + "submission_ms": 0.082528, + "synchronized_e2e_ms": 0.196641 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.05408, + "submission_ms": 0.05408, + "synchronized_e2e_ms": 0.168288 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.210304, + "submission_ms": 0.210304, + "synchronized_e2e_ms": 0.318016 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.792904, + "submission_ms": 7.792904, + "synchronized_e2e_ms": 7.882312 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.105344, + "submission_ms": 0.105344, + "synchronized_e2e_ms": 0.186048 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.053888, + "submission_ms": 0.053888, + "synchronized_e2e_ms": 0.14368 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.358017, + "submission_ms": 1.358017, + "synchronized_e2e_ms": 1.382657 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.044193, + "submission_ms": 1.044193, + "synchronized_e2e_ms": 1.069665 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.09328, + "submission_ms": 0.09328, + "synchronized_e2e_ms": 0.191808 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 79.24757, + "submission_ms": 79.24757, + "synchronized_e2e_ms": 79.325138 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 2.001693, + "evolution_kernel_ms": 0.302528, + "evolution_speedup": 6.6166, + "evolution_tflops": 3.9929, + "expected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d144_b1_n512_k8192_d144", + "measurement_order": [ + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 3360, + "measurement_schedule_sha256": "dc0691e40c3d52a5c002d72d10d7f8259e5f86ba41d7a53da2e732ba53a4756c", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 791, + "public_pair_count": 889, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 791, + "baseline_public_raw": 889, + "candidate_precomputed": 791, + "candidate_public_raw": 889 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 674 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:post_d895_d144_b1_n512_k8192_d144", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.1299245599329422, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.6263119797321752, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.0993671388255259, + "including_init_synchronized_e2e_speedup": 0.07723037119783416, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.12939946359107954, + "including_init_synchronized_e2e_speedup": 0.08159368326365933, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.3754266509084347, + "including_init_synchronized_e2e_speedup": 0.12391192053329766, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.1531306162279533, + "including_init_synchronized_e2e_speedup": 0.44699185630097643, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.5160645780051152, + "hot_synchronized_e2e_speedup": 1.6263119797321752, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 214404, + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d144_b1_n512_k8192_d144", + "source": "high_k_low_n", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 128, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 8 + } + }, + { + "B": 2, + "D": 144, + "K": 1024, + "N": 2048, + "baseline_07cf_adapter_bench_iters": 3397, + "baseline_07cf_adapter_gpu_span_ms": 0.069761, + "baseline_07cf_adapter_host_enqueue_ms": 0.152128, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.036672, + "baseline_07cf_adapter_kernel_sum_ms": 0.033088, + "baseline_07cf_adapter_submission_ms": 0.152128, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.17472, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.033088 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.201696, + "synchronized_e2e_ms": 0.22144 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.069761 + }, + "host_enqueue_ms": { + "median": 0.152128 + }, + "inter_kernel_gap_ms": { + "median": 0.036672 + }, + "kernel_sum_ms": { + "median": 0.033088 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3397, + "submission_ms": { + "median": 0.152128 + }, + "synchronized_e2e_ms": { + "median": 0.17472 + } + }, + "baseline_07cf_precomputed_bench_iters": 4080, + "baseline_07cf_precomputed_gpu_span_ms": 0.028096, + "baseline_07cf_precomputed_host_enqueue_ms": 0.044736, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.028096, + "baseline_07cf_precomputed_submission_ms": 0.044736, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.078528, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.028096 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.039137, + "synchronized_e2e_ms": 0.065505 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.028096 + }, + "host_enqueue_ms": { + "median": 0.044736 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.028096 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4080, + "submission_ms": { + "median": 0.044736 + }, + "synchronized_e2e_ms": { + "median": 0.078528 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.4829513097949887, + "submission": 3.4005722460658085, + "synchronized_e2e": 2.2249388753056234 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.60868, + "after_init_synchronized_e2e_ms_per_call": 7.635976, + "including_init_host_enqueue_ms_per_call": 44.543118, + "including_init_synchronized_e2e_ms_per_call": 462.74729500000007, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8977831999999999, + "after_init_synchronized_e2e_ms_per_call": 0.9208456, + "including_init_host_enqueue_ms_per_call": 4.591227, + "including_init_synchronized_e2e_ms_per_call": 46.4319775, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22669352000000004, + "after_init_synchronized_e2e_ms_per_call": 0.24933255999999995, + "including_init_host_enqueue_ms_per_call": 0.5960379, + "including_init_synchronized_e2e_ms_per_call": 4.800445750000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.159584552, + "after_init_synchronized_e2e_ms_per_call": 0.18218125599999999, + "including_init_host_enqueue_ms_per_call": 0.19651899, + "including_init_synchronized_e2e_ms_per_call": 0.637292575, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.60868, + "synchronized_e2e_ms": 7.635976, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.402688, + "median": 0.069761, + "min": 0.058752, + "p90": 0.09814400000000004, + "sample_count": 3397 + }, + "host_enqueue_ms": { + "max": 44.48619, + "median": 0.152128, + "min": 0.121696, + "p90": 0.2503424, + "sample_count": 3397 + }, + "sample_count": 3397, + "synchronized_e2e_ms": { + "max": 44.696718, + "median": 0.17472, + "min": 0.144256, + "p90": 0.28012200000000015, + "sample_count": 3397 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 4080, + "candidate_precomputed_gpu_span_ms": 0.025184, + "candidate_precomputed_host_enqueue_ms": 0.055616, + "candidate_precomputed_inter_kernel_gap_ms": 0.005152, + "candidate_precomputed_kernel_sum_ms": 0.020032, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.055616, + "candidate_precomputed_synchronized_e2e_ms": 0.0688, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.020032 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.037248, + "synchronized_e2e_ms": 0.05024 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.025184 + }, + "host_enqueue_ms": { + "median": 0.055616 + }, + "inter_kernel_gap_ms": { + "median": 0.005152 + }, + "kernel_sum_ms": { + "median": 0.020032 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4080, + "submission_ms": { + "median": 0.055616 + }, + "synchronized_e2e_ms": { + "median": 0.0688 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01b4a090", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01b481a0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.1715374841168995, + "submission": 0.833716915995397, + "synchronized_e2e": 1.1218604651162791 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 3397, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.029504, + "candidate_public_raw_host_enqueue_ms": 0.046368, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.029696, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.046368, + "candidate_public_raw_synchronized_e2e_ms": 0.077184, + "candidate_public_raw_tflops_from_gpu_span": 40.94222993492408, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.029504 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.050016, + "synchronized_e2e_ms": 0.072192 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.029504 + }, + "host_enqueue_ms": { + "median": 0.046368 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.029696 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3397, + "submission_ms": { + "median": 0.046368 + }, + "synchronized_e2e_ms": { + "median": 0.077184 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 75.600879, + "after_init_synchronized_e2e_ms_per_call": 75.647055, + "including_init_host_enqueue_ms_per_call": 112.958902, + "including_init_synchronized_e2e_ms_per_call": 113.08267799999999, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 7.6018191, + "after_init_synchronized_e2e_ms_per_call": 7.634171099999999, + "including_init_host_enqueue_ms_per_call": 11.3376214, + "including_init_synchronized_e2e_ms_per_call": 11.377733399999999, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8019131100000001, + "after_init_synchronized_e2e_ms_per_call": 0.8328827099999999, + "including_init_host_enqueue_ms_per_call": 1.1754933399999998, + "including_init_synchronized_e2e_ms_per_call": 1.2072389399999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.12192251100000001, + "after_init_synchronized_e2e_ms_per_call": 0.152753871, + "including_init_host_enqueue_ms_per_call": 0.159280534, + "including_init_synchronized_e2e_ms_per_call": 0.190189494, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 75.600879, + "synchronized_e2e_ms": 75.647055, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.030336, + "median": 0.029504, + "min": 0.029152, + "p90": 0.029888, + "sample_count": 3397 + }, + "host_enqueue_ms": { + "max": 50.732149, + "median": 0.046368, + "min": 0.035904, + "p90": 0.07966120000000002, + "sample_count": 3397 + }, + "sample_count": 3397, + "synchronized_e2e_ms": { + "max": 50.865621, + "median": 0.077184, + "min": 0.067616, + "p90": 0.10659240000000002, + "sample_count": 3397 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.064384, + "submission_ms": 0.064384, + "synchronized_e2e_ms": 0.089632 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.039137, + "submission_ms": 0.039137, + "synchronized_e2e_ms": 0.065505 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.201696, + "submission_ms": 0.201696, + "synchronized_e2e_ms": 0.22144 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.60868, + "submission_ms": 7.60868, + "synchronized_e2e_ms": 7.635976 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.083584, + "submission_ms": 0.083584, + "synchronized_e2e_ms": 0.100768 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.037248, + "submission_ms": 0.037248, + "synchronized_e2e_ms": 0.05024 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.181441, + "submission_ms": 1.181441, + "synchronized_e2e_ms": 1.206369 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.016609, + "submission_ms": 1.016609, + "synchronized_e2e_ms": 1.038945 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.050016, + "submission_ms": 0.050016, + "synchronized_e2e_ms": 0.072192 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 75.600879, + "submission_ms": 75.600879, + "synchronized_e2e_ms": 75.647055 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.465663, + "evolution_kernel_ms": 0.192896, + "evolution_speedup": 2.4141, + "evolution_tflops": 6.2622, + "expected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d144_b2_n2048_k1024_d144", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 14954, + "measurement_schedule_sha256": "7e57c23f0a6c2285c6ee45fe0245ed09de9da0fc2076464baa424c19c4837a50", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 4080, + "public_pair_count": 3397, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 4080, + "baseline_public_raw": 3397, + "candidate_precomputed": 4080, + "candidate_public_raw": 3397 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2992 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:post_d895_d144_b2_n2048_k1024_d144", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.1156289707750953, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.263681592039801, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.10094214507094824, + "including_init_synchronized_e2e_speedup": 4.092114753419619, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.12062155641232618, + "including_init_synchronized_e2e_speedup": 4.080951439765675, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.29936095083544234, + "including_init_synchronized_e2e_speedup": 3.976384119948948, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.1926457562571358, + "including_init_synchronized_e2e_speedup": 3.3508295416149543, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.3644590563991326, + "hot_synchronized_e2e_speedup": 2.263681592039801, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 214402, + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d144_b2_n2048_k1024_d144", + "source": "tailpad_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 144, + "K": 1024, + "N": 8192, + "baseline_07cf_adapter_bench_iters": 1525, + "baseline_07cf_adapter_gpu_span_ms": 0.09584, + "baseline_07cf_adapter_host_enqueue_ms": 0.146432, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.002656, + "baseline_07cf_adapter_kernel_sum_ms": 0.093088, + "baseline_07cf_adapter_submission_ms": 0.146432, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.195584, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.093088 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.291232, + "synchronized_e2e_ms": 0.333088 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.09584 + }, + "host_enqueue_ms": { + "median": 0.146432 + }, + "inter_kernel_gap_ms": { + "median": 0.002656 + }, + "kernel_sum_ms": { + "median": 0.093088 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1525, + "submission_ms": { + "median": 0.146432 + }, + "synchronized_e2e_ms": { + "median": 0.195584 + } + }, + "baseline_07cf_precomputed_bench_iters": 3446, + "baseline_07cf_precomputed_gpu_span_ms": 0.05296, + "baseline_07cf_precomputed_host_enqueue_ms": 0.041152, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.05296, + "baseline_07cf_precomputed_submission_ms": 0.041152, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.100736, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.05296 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.049216, + "synchronized_e2e_ms": 0.104384 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.05296 + }, + "host_enqueue_ms": { + "median": 0.041152 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.05296 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3446, + "submission_ms": { + "median": 0.041152 + }, + "synchronized_e2e_ms": { + "median": 0.100736 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.809667673716012, + "submission": 3.558320373250389, + "synchronized_e2e": 1.9415501905972046 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.029193, + "after_init_synchronized_e2e_ms_per_call": 8.055817, + "including_init_host_enqueue_ms_per_call": 42.2571, + "including_init_synchronized_e2e_ms_per_call": 42.40126, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9347080999999999, + "after_init_synchronized_e2e_ms_per_call": 0.9816073, + "including_init_host_enqueue_ms_per_call": 4.3574988, + "including_init_synchronized_e2e_ms_per_call": 4.4161516, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22525961000000003, + "after_init_synchronized_e2e_ms_per_call": 0.27418633, + "including_init_host_enqueue_ms_per_call": 0.5675386800000001, + "including_init_synchronized_e2e_ms_per_call": 0.61764076, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.154314761, + "after_init_synchronized_e2e_ms_per_call": 0.203444233, + "including_init_host_enqueue_ms_per_call": 0.18854266800000002, + "including_init_synchronized_e2e_ms_per_call": 0.237789676, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.029193, + "synchronized_e2e_ms": 8.055817, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.302272, + "median": 0.09584, + "min": 0.092544, + "p90": 0.11063040000000005, + "sample_count": 1525 + }, + "host_enqueue_ms": { + "max": 0.986721, + "median": 0.146432, + "min": 0.124704, + "p90": 0.18185640000000003, + "sample_count": 1525 + }, + "sample_count": 1525, + "synchronized_e2e_ms": { + "max": 1.055969, + "median": 0.195584, + "min": 0.179104, + "p90": 0.2273152, + "sample_count": 1525 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3446, + "candidate_precomputed_gpu_span_ms": 0.026432, + "candidate_precomputed_host_enqueue_ms": 0.051744, + "candidate_precomputed_inter_kernel_gap_ms": 0.001792, + "candidate_precomputed_kernel_sum_ms": 0.02464, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.051744, + "candidate_precomputed_synchronized_e2e_ms": 0.064, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.02464 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.041408, + "synchronized_e2e_ms": 0.05616 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.026432 + }, + "host_enqueue_ms": { + "median": 0.051744 + }, + "inter_kernel_gap_ms": { + "median": 0.001792 + }, + "kernel_sum_ms": { + "median": 0.02464 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3446, + "submission_ms": { + "median": 0.051744 + }, + "synchronized_e2e_ms": { + "median": 0.064 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295c2f30", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295c2ea0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.4854721549636807, + "submission": 0.8509778911564627, + "synchronized_e2e": 1.7545 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 1525, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.065696, + "candidate_public_raw_host_enqueue_ms": 0.044033, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.065888, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.044033, + "candidate_public_raw_synchronized_e2e_ms": 0.112288, + "candidate_public_raw_tflops_from_gpu_span": 147.0968767657087, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.065696 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.060928, + "synchronized_e2e_ms": 0.108768 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.065696 + }, + "host_enqueue_ms": { + "median": 0.044033 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.065888 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1525, + "submission_ms": { + "median": 0.044033 + }, + "synchronized_e2e_ms": { + "median": 0.112288 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 80.869556, + "after_init_synchronized_e2e_ms_per_call": 80.918292, + "including_init_host_enqueue_ms_per_call": 115.447096, + "including_init_synchronized_e2e_ms_per_call": 529.006084, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 8.1265853, + "after_init_synchronized_e2e_ms_per_call": 8.1928884, + "including_init_host_enqueue_ms_per_call": 11.5843393, + "including_init_synchronized_e2e_ms_per_call": 53.0016676, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8522882300000001, + "after_init_synchronized_e2e_ms_per_call": 0.9203480399999999, + "including_init_host_enqueue_ms_per_call": 1.19806363, + "including_init_synchronized_e2e_ms_per_call": 5.40122596, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.124858523, + "after_init_synchronized_e2e_ms_per_call": 0.19309400399999999, + "including_init_host_enqueue_ms_per_call": 0.159436063, + "including_init_synchronized_e2e_ms_per_call": 0.6411817959999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 80.869556, + "synchronized_e2e_ms": 80.918292, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.067104, + "median": 0.065696, + "min": 0.065088, + "p90": 0.066304, + "sample_count": 1525 + }, + "host_enqueue_ms": { + "max": 26.824604, + "median": 0.044033, + "min": 0.037248, + "p90": 0.05770280000000003, + "sample_count": 1525 + }, + "sample_count": 1525, + "synchronized_e2e_ms": { + "max": 33.882691, + "median": 0.112288, + "min": 0.106464, + "p90": 0.12655360000000004, + "sample_count": 1525 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.07248, + "submission_ms": 0.07248, + "synchronized_e2e_ms": 0.126912 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.049216, + "submission_ms": 0.049216, + "synchronized_e2e_ms": 0.104384 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.291232, + "submission_ms": 0.291232, + "synchronized_e2e_ms": 0.333088 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.029193, + "submission_ms": 8.029193, + "synchronized_e2e_ms": 8.055817 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.069472, + "submission_ms": 0.069472, + "synchronized_e2e_ms": 0.087008 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.041408, + "submission_ms": 0.041408, + "synchronized_e2e_ms": 0.05616 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.424834, + "submission_ms": 1.424834, + "synchronized_e2e_ms": 1.449026 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.168993, + "submission_ms": 1.168993, + "synchronized_e2e_ms": 1.190785 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.060928, + "submission_ms": 0.060928, + "synchronized_e2e_ms": 0.108768 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 80.869556, + "submission_ms": 80.869556, + "synchronized_e2e_ms": 80.918292 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.473791, + "evolution_kernel_ms": 0.195104, + "evolution_speedup": 2.4284, + "evolution_tflops": 49.5309, + "expected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d144_b4_n8192_k1024_d144", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 9942, + "measurement_schedule_sha256": "cb11b45182ab0adeb1e06363c4eec185c68e6cff94813229a592045d7537e522", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3446, + "public_pair_count": 1525, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3446, + "baseline_public_raw": 1525, + "candidate_precomputed": 3446, + "candidate_public_raw": 1525 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 1990 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:post_d895_d144_b4_n8192_k1024_d144", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 2.00363196125908, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.7418067825591337, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.09955495600426167, + "including_init_synchronized_e2e_speedup": 0.08015268875433199, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.11981211656685084, + "including_init_synchronized_e2e_speedup": 0.08332099347002433, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.29791591667865125, + "including_init_synchronized_e2e_speedup": 0.1143519572360198, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.0536020217385933, + "including_init_synchronized_e2e_speedup": 0.3708615520332084, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.4588407208962493, + "hot_synchronized_e2e_speedup": 1.7418067825591337, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 214403, + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d144_b4_n8192_k1024_d144", + "source": "tailpad_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 16, + "K": 1024, + "N": 32768, + "baseline_07cf_adapter_bench_iters": 785, + "baseline_07cf_adapter_gpu_span_ms": 0.141504, + "baseline_07cf_adapter_host_enqueue_ms": 0.152832, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.002112, + "baseline_07cf_adapter_kernel_sum_ms": 0.13936, + "baseline_07cf_adapter_submission_ms": 0.152832, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.245089, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.13936 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.193057, + "synchronized_e2e_ms": 0.256737 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.141504 + }, + "host_enqueue_ms": { + "median": 0.152832 + }, + "inter_kernel_gap_ms": { + "median": 0.002112 + }, + "kernel_sum_ms": { + "median": 0.13936 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 785, + "submission_ms": { + "median": 0.152832 + }, + "synchronized_e2e_ms": { + "median": 0.245089 + } + }, + "baseline_07cf_precomputed_bench_iters": 1609, + "baseline_07cf_precomputed_gpu_span_ms": 0.067904, + "baseline_07cf_precomputed_host_enqueue_ms": 0.04272, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.067904, + "baseline_07cf_precomputed_submission_ms": 0.04272, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.118112, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.067904 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.042368, + "synchronized_e2e_ms": 0.11504 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.067904 + }, + "host_enqueue_ms": { + "median": 0.04272 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.067904 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1609, + "submission_ms": { + "median": 0.04272 + }, + "synchronized_e2e_ms": { + "median": 0.118112 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.083883129123468, + "submission": 3.57752808988764, + "synchronized_e2e": 2.075055879165538 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.467337, + "after_init_synchronized_e2e_ms_per_call": 8.535017, + "including_init_host_enqueue_ms_per_call": 42.916781, + "including_init_synchronized_e2e_ms_per_call": 492.576511, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9842825000000002, + "after_init_synchronized_e2e_ms_per_call": 1.0740818, + "including_init_host_enqueue_ms_per_call": 4.4292269, + "including_init_synchronized_e2e_ms_per_call": 49.478231199999996, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23597704999999997, + "after_init_synchronized_e2e_ms_per_call": 0.32798828, + "including_init_host_enqueue_ms_per_call": 0.58047149, + "including_init_synchronized_e2e_ms_per_call": 5.16840322, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.161146505, + "after_init_synchronized_e2e_ms_per_call": 0.253378928, + "including_init_host_enqueue_ms_per_call": 0.19559594900000002, + "including_init_synchronized_e2e_ms_per_call": 0.7374204219999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.467337, + "synchronized_e2e_ms": 8.535017, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.418241, + "median": 0.141504, + "min": 0.140064, + "p90": 0.14233659999999998, + "sample_count": 785 + }, + "host_enqueue_ms": { + "max": 0.757952, + "median": 0.152832, + "min": 0.132096, + "p90": 0.1802688, + "sample_count": 785 + }, + "sample_count": 785, + "synchronized_e2e_ms": { + "max": 0.814528, + "median": 0.245089, + "min": 0.2328, + "p90": 0.2657984, + "sample_count": 785 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 1609, + "candidate_precomputed_gpu_span_ms": 0.059136, + "candidate_precomputed_host_enqueue_ms": 0.054497, + "candidate_precomputed_inter_kernel_gap_ms": 0.002656, + "candidate_precomputed_kernel_sum_ms": 0.05648, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.054497, + "candidate_precomputed_synchronized_e2e_ms": 0.09472, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.05648 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.043328, + "synchronized_e2e_ms": 0.0872 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.059136 + }, + "host_enqueue_ms": { + "median": 0.054497 + }, + "inter_kernel_gap_ms": { + "median": 0.002656 + }, + "kernel_sum_ms": { + "median": 0.05648 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1609, + "submission_ms": { + "median": 0.054497 + }, + "synchronized_e2e_ms": { + "median": 0.09472 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc045e8230", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc045e9370" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.152056277056277, + "submission": 0.8608180266803678, + "synchronized_e2e": 1.8689189189189188 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 785, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.127264, + "candidate_public_raw_host_enqueue_ms": 0.046912, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.127488, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.046912, + "candidate_public_raw_synchronized_e2e_ms": 0.177024, + "candidate_public_raw_tflops_from_gpu_span": 33.748485793311545, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.127264 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.081248, + "synchronized_e2e_ms": 0.196992 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.127264 + }, + "host_enqueue_ms": { + "median": 0.046912 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.127488 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 785, + "submission_ms": { + "median": 0.046912 + }, + "synchronized_e2e_ms": { + "median": 0.177024 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 72.919883, + "after_init_synchronized_e2e_ms_per_call": 73.013579, + "including_init_host_enqueue_ms_per_call": 107.793967, + "including_init_synchronized_e2e_ms_per_call": 107.95665499999998, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 7.3342091, + "after_init_synchronized_e2e_ms_per_call": 7.4606794999999995, + "including_init_host_enqueue_ms_per_call": 10.821617499999999, + "including_init_synchronized_e2e_ms_per_call": 10.954987099999999, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.77564171, + "after_init_synchronized_e2e_ms_per_call": 0.9053895499999999, + "including_init_host_enqueue_ms_per_call": 1.12438255, + "including_init_synchronized_e2e_ms_per_call": 1.2548203099999997, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.119784971, + "after_init_synchronized_e2e_ms_per_call": 0.24986055499999998, + "including_init_host_enqueue_ms_per_call": 0.154659055, + "including_init_synchronized_e2e_ms_per_call": 0.284803631, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 72.919883, + "synchronized_e2e_ms": 73.013579, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.127872, + "median": 0.127264, + "min": 0.126881, + "p90": 0.12752, + "sample_count": 785 + }, + "host_enqueue_ms": { + "max": 0.089312, + "median": 0.046912, + "min": 0.03856, + "p90": 0.0557312, + "sample_count": 785 + }, + "sample_count": 785, + "synchronized_e2e_ms": { + "max": 0.219136, + "median": 0.177024, + "min": 0.169568, + "p90": 0.18433919999999998, + "sample_count": 785 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.077248, + "submission_ms": 0.077248, + "synchronized_e2e_ms": 0.14864 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.042368, + "submission_ms": 0.042368, + "synchronized_e2e_ms": 0.11504 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.193057, + "submission_ms": 0.193057, + "synchronized_e2e_ms": 0.256737 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.467337, + "submission_ms": 8.467337, + "synchronized_e2e_ms": 8.535017 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.081152, + "submission_ms": 0.081152, + "synchronized_e2e_ms": 0.113504 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.043328, + "submission_ms": 0.043328, + "synchronized_e2e_ms": 0.0872 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.475233, + "submission_ms": 1.475233, + "synchronized_e2e_ms": 1.496641 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.176354, + "submission_ms": 1.176354, + "synchronized_e2e_ms": 1.19853 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.081248, + "submission_ms": 0.081248, + "synchronized_e2e_ms": 0.196992 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 72.919883, + "submission_ms": 72.919883, + "synchronized_e2e_ms": 73.013579 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.330559, + "evolution_kernel_ms": 0.216591, + "evolution_speedup": 1.5262, + "evolution_tflops": 19.8298, + "expected_route": "microdim_hybrid_9c0d_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d16_b4_n32768_k1024_d16", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 4788, + "measurement_schedule_sha256": "350515cf1eae283a6d09d9d480f5401a4b6961da9d18ec60df87f7440ebc2015", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 1609, + "public_pair_count": 785, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 1609, + "baseline_public_raw": 785, + "candidate_precomputed": 1609, + "candidate_public_raw": 785 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 958 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:post_d895_d16_b4_n32768_k1024_d16", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.1482683982683983, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.3844958875632685, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.11689629678336957, + "including_init_synchronized_e2e_speedup": 4.562724836185412, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.14396568033782983, + "including_init_synchronized_e2e_speedup": 4.516502917652911, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.36226205615030577, + "including_init_synchronized_e2e_speedup": 4.118839310147922, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.0140813462933356, + "including_init_synchronized_e2e_speedup": 2.5892240889302425, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.1118933869751069, + "hot_synchronized_e2e_speedup": 1.3844958875632685, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 21602, + "selected_route": "microdim_hybrid_9c0d_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d16_b4_n32768_k1024_d16", + "source": "near_floor_microdim", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 16, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 8, + "D": 16, + "K": 512, + "N": 65536, + "baseline_07cf_adapter_bench_iters": 290, + "baseline_07cf_adapter_gpu_span_ms": 0.3692795, + "baseline_07cf_adapter_host_enqueue_ms": 0.17912, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.00208, + "baseline_07cf_adapter_kernel_sum_ms": 0.3672, + "baseline_07cf_adapter_submission_ms": 0.17912, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.487552, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.3672 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.259072, + "synchronized_e2e_ms": 0.520961 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.3692795 + }, + "host_enqueue_ms": { + "median": 0.17912 + }, + "inter_kernel_gap_ms": { + "median": 0.00208 + }, + "kernel_sum_ms": { + "median": 0.3672 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 290, + "submission_ms": { + "median": 0.17912 + }, + "synchronized_e2e_ms": { + "median": 0.487552 + } + }, + "baseline_07cf_precomputed_bench_iters": 1496, + "baseline_07cf_precomputed_gpu_span_ms": 0.091936, + "baseline_07cf_precomputed_host_enqueue_ms": 0.050992499999999996, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.091936, + "baseline_07cf_precomputed_submission_ms": 0.050992499999999996, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.1496, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.091936 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.068256, + "synchronized_e2e_ms": 0.16224 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.091936 + }, + "host_enqueue_ms": { + "median": 0.050992499999999996 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.091936 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1496, + "submission_ms": { + "median": 0.050992499999999996 + }, + "synchronized_e2e_ms": { + "median": 0.1496 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 4.016701836059867, + "submission": 3.5126734323675053, + "synchronized_e2e": 3.25903743315508 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.059784, + "after_init_synchronized_e2e_ms_per_call": 8.127304, + "including_init_host_enqueue_ms_per_call": 43.611052, + "including_init_synchronized_e2e_ms_per_call": 43.758541, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9671864000000001, + "after_init_synchronized_e2e_ms_per_call": 1.2515272, + "including_init_host_enqueue_ms_per_call": 4.5223132, + "including_init_synchronized_e2e_ms_per_call": 4.8146509, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.25792664000000004, + "after_init_synchronized_e2e_ms_per_call": 0.56394952, + "including_init_host_enqueue_ms_per_call": 0.61343932, + "including_init_synchronized_e2e_ms_per_call": 0.92026189, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.187000664, + "after_init_synchronized_e2e_ms_per_call": 0.49519175199999993, + "including_init_host_enqueue_ms_per_call": 0.22255193199999998, + "including_init_synchronized_e2e_ms_per_call": 0.530822989, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.059784, + "synchronized_e2e_ms": 8.127304, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.510017, + "median": 0.3692795, + "min": 0.367776, + "p90": 0.370208, + "sample_count": 290 + }, + "host_enqueue_ms": { + "max": 0.562528, + "median": 0.17912, + "min": 0.14992, + "p90": 0.2120352, + "sample_count": 290 + }, + "sample_count": 290, + "synchronized_e2e_ms": { + "max": 0.630784, + "median": 0.487552, + "min": 0.473537, + "p90": 0.5086858000000001, + "sample_count": 290 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 1496, + "candidate_precomputed_gpu_span_ms": 0.066816, + "candidate_precomputed_host_enqueue_ms": 0.045424, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.066816, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.045424, + "candidate_precomputed_synchronized_e2e_ms": 0.107648, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.066816 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.047488, + "synchronized_e2e_ms": 0.114656 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.066816 + }, + "host_enqueue_ms": { + "median": 0.045424 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.066816 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1496, + "submission_ms": { + "median": 0.045424 + }, + "synchronized_e2e_ms": { + "median": 0.107648 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb937d1e50", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb937d1760" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 5.1767316211685825, + "submission": 1.3571680169073619, + "synchronized_e2e": 3.7990533962544593 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 290, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.3458885, + "candidate_public_raw_host_enqueue_ms": 0.061648, + "candidate_public_raw_inter_kernel_gap_ms": 0.00016, + "candidate_public_raw_kernel_sum_ms": 0.345728, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.061648, + "candidate_public_raw_synchronized_e2e_ms": 0.4089605, + "candidate_public_raw_tflops_from_gpu_span": 24.834403549120598, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.345728 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.293856, + "synchronized_e2e_ms": 0.633568 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.3458885 + }, + "host_enqueue_ms": { + "median": 0.061648 + }, + "inter_kernel_gap_ms": { + "median": 0.00016 + }, + "kernel_sum_ms": { + "median": 0.345728 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 290, + "submission_ms": { + "median": 0.061648 + }, + "synchronized_e2e_ms": { + "median": 0.4089605 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 40.892715, + "after_init_synchronized_e2e_ms_per_call": 41.223083, + "including_init_host_enqueue_ms_per_call": 76.73073500000001, + "including_init_synchronized_e2e_ms_per_call": 525.323296, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 4.1447547, + "after_init_synchronized_e2e_ms_per_call": 4.490372750000001, + "including_init_host_enqueue_ms_per_call": 7.728556700000001, + "including_init_synchronized_e2e_ms_per_call": 52.90039405, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.46995867, + "after_init_synchronized_e2e_ms_per_call": 0.817101725, + "including_init_host_enqueue_ms_per_call": 0.8283388700000001, + "including_init_synchronized_e2e_ms_per_call": 5.658103855, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.10247906700000001, + "after_init_synchronized_e2e_ms_per_call": 0.44977462249999994, + "including_init_host_enqueue_ms_per_call": 0.138317087, + "including_init_synchronized_e2e_ms_per_call": 0.9338748355000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 40.892715, + "synchronized_e2e_ms": 41.223083, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.346817, + "median": 0.3458885, + "min": 0.345248, + "p90": 0.3461441, + "sample_count": 290 + }, + "host_enqueue_ms": { + "max": 0.0976, + "median": 0.061648, + "min": 0.053184, + "p90": 0.0723808, + "sample_count": 290 + }, + "sample_count": 290, + "synchronized_e2e_ms": { + "max": 0.44048, + "median": 0.4089605, + "min": 0.401376, + "p90": 0.4189888, + "sample_count": 290 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.089504, + "submission_ms": 0.089504, + "synchronized_e2e_ms": 0.181856 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.068256, + "submission_ms": 0.068256, + "synchronized_e2e_ms": 0.16224 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.259072, + "submission_ms": 0.259072, + "synchronized_e2e_ms": 0.520961 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.059784, + "submission_ms": 8.059784, + "synchronized_e2e_ms": 8.127304 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.07488, + "submission_ms": 0.07488, + "synchronized_e2e_ms": 0.129728 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.047488, + "submission_ms": 0.047488, + "synchronized_e2e_ms": 0.114656 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.296898, + "submission_ms": 1.296898, + "synchronized_e2e_ms": 1.321186 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.107425, + "submission_ms": 1.107425, + "synchronized_e2e_ms": 1.130465 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.293856, + "submission_ms": 0.293856, + "synchronized_e2e_ms": 0.633568 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 40.892715, + "submission_ms": 40.892715, + "synchronized_e2e_ms": 41.223083 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.444799, + "evolution_kernel_ms": 0.295359, + "evolution_speedup": 1.506, + "evolution_tflops": 29.083, + "expected_route": "microdim_pipeline4_08f9_v4", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d16_b8_n65536_k512_d16", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 3572, + "measurement_schedule_sha256": "293f26dc1583a088345561a616724fc3c684a4e36bac444f768e543fbeb45c5d", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 1496, + "public_pair_count": 290, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 1496, + "baseline_public_raw": 290, + "candidate_precomputed": 1496, + "candidate_public_raw": 290 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 716 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:post_d895_d16_b8_n65536_k512_d16", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.3759578544061304, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.192173816297662, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.19715420120324334, + "including_init_synchronized_e2e_speedup": 0.08329830664886409, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.2787134319751071, + "including_init_synchronized_e2e_speedup": 0.09101351675092106, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.6901827553968265, + "including_init_synchronized_e2e_speedup": 0.16264492727307847, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.1009775279173293, + "including_init_synchronized_e2e_speedup": 0.5684091366652957, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.067625838962556, + "hot_synchronized_e2e_speedup": 1.192173816297662, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 21601, + "selected_route": "microdim_pipeline4_08f9_v4", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d16_b8_n65536_k512_d16", + "source": "near_floor_microdim", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 128, + "D_PAD": 16, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 176, + "K": 256, + "N": 256, + "baseline_07cf_adapter_bench_iters": 9616, + "baseline_07cf_adapter_gpu_span_ms": 0.054976, + "baseline_07cf_adapter_host_enqueue_ms": 0.1592325, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.044704, + "baseline_07cf_adapter_kernel_sum_ms": 0.010304, + "baseline_07cf_adapter_submission_ms": 0.1592325, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.1820485, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.010304 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.241056, + "synchronized_e2e_ms": 0.264064 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.054976 + }, + "host_enqueue_ms": { + "median": 0.1592325 + }, + "inter_kernel_gap_ms": { + "median": 0.044704 + }, + "kernel_sum_ms": { + "median": 0.010304 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 9616, + "submission_ms": { + "median": 0.1592325 + }, + "synchronized_e2e_ms": { + "median": 0.1820485 + } + }, + "baseline_07cf_precomputed_bench_iters": 11242, + "baseline_07cf_precomputed_gpu_span_ms": 0.008896, + "baseline_07cf_precomputed_host_enqueue_ms": 0.046848, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.008896, + "baseline_07cf_precomputed_submission_ms": 0.046848, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.062368, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.008896 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.047424, + "synchronized_e2e_ms": 0.063936 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.008896 + }, + "host_enqueue_ms": { + "median": 0.046848 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.008896 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 11242, + "submission_ms": { + "median": 0.046848 + }, + "synchronized_e2e_ms": { + "median": 0.062368 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 6.179856115107913, + "submission": 3.3989177766393444, + "synchronized_e2e": 2.918940802975885 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.947785, + "after_init_synchronized_e2e_ms_per_call": 8.987498, + "including_init_host_enqueue_ms_per_call": 45.882222999999996, + "including_init_synchronized_e2e_ms_per_call": 464.09881700000005, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 1.03808775, + "after_init_synchronized_e2e_ms_per_call": 1.06259345, + "including_init_host_enqueue_ms_per_call": 4.73153155, + "including_init_synchronized_e2e_ms_per_call": 46.573725350000004, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.24711802499999996, + "after_init_synchronized_e2e_ms_per_call": 0.270102995, + "including_init_host_enqueue_ms_per_call": 0.6164624049999999, + "including_init_synchronized_e2e_ms_per_call": 4.821216185000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.1680210525, + "after_init_synchronized_e2e_ms_per_call": 0.1908539495, + "including_init_host_enqueue_ms_per_call": 0.2049554905, + "including_init_synchronized_e2e_ms_per_call": 0.6459652685000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.947785, + "synchronized_e2e_ms": 8.987498, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 84.544215, + "median": 0.054976, + "min": 0.043104, + "p90": 0.089104, + "sample_count": 9616 + }, + "host_enqueue_ms": { + "max": 84.94364, + "median": 0.1592325, + "min": 0.125216, + "p90": 0.331504, + "sample_count": 9616 + }, + "sample_count": 9616, + "synchronized_e2e_ms": { + "max": 85.219161, + "median": 0.1820485, + "min": 0.144736, + "p90": 0.3764965, + "sample_count": 9616 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 11242, + "candidate_precomputed_gpu_span_ms": 0.0144, + "candidate_precomputed_host_enqueue_ms": 0.057824, + "candidate_precomputed_inter_kernel_gap_ms": 0.006272, + "candidate_precomputed_kernel_sum_ms": 0.008192, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.057824, + "candidate_precomputed_synchronized_e2e_ms": 0.071552, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.008192 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.044384, + "synchronized_e2e_ms": 0.062368 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.0144 + }, + "host_enqueue_ms": { + "median": 0.057824 + }, + "inter_kernel_gap_ms": { + "median": 0.006272 + }, + "kernel_sum_ms": { + "median": 0.008192 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 11242, + "submission_ms": { + "median": 0.057824 + }, + "synchronized_e2e_ms": { + "median": 0.071552 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc295c2cc0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc2819c230" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 0.7666666666666666, + "submission": 0.8489208633093525, + "synchronized_e2e": 0.9101073345259391 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 9616, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.01104, + "candidate_public_raw_host_enqueue_ms": 0.049088, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.011232, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.049088, + "candidate_public_raw_synchronized_e2e_ms": 0.06512, + "candidate_public_raw_tflops_from_gpu_span": 2.089553623188406, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.01104 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.060224, + "synchronized_e2e_ms": 0.082336 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.01104 + }, + "host_enqueue_ms": { + "median": 0.049088 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.011232 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 9616, + "submission_ms": { + "median": 0.049088 + }, + "synchronized_e2e_ms": { + "median": 0.06512 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 41.918091, + "after_init_synchronized_e2e_ms_per_call": 41.951979, + "including_init_host_enqueue_ms_per_call": 79.27611399999999, + "including_init_synchronized_e2e_ms_per_call": 79.387602, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 4.2359883, + "after_init_synchronized_e2e_ms_per_call": 4.253805900000001, + "including_init_host_enqueue_ms_per_call": 7.9717906, + "including_init_synchronized_e2e_ms_per_call": 7.9973681999999995, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.46777803, + "after_init_synchronized_e2e_ms_per_call": 0.48398859, + "including_init_host_enqueue_ms_per_call": 0.8413582599999999, + "including_init_synchronized_e2e_ms_per_call": 0.8583448200000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.09095700299999998, + "after_init_synchronized_e2e_ms_per_call": 0.107006859, + "including_init_host_enqueue_ms_per_call": 0.128315026, + "including_init_synchronized_e2e_ms_per_call": 0.14444248199999998, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 41.918091, + "synchronized_e2e_ms": 41.951979, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.011648, + "median": 0.01104, + "min": 0.010272, + "p90": 0.011264, + "sample_count": 9616 + }, + "host_enqueue_ms": { + "max": 56.904251, + "median": 0.049088, + "min": 0.03616, + "p90": 0.10513600000000001, + "sample_count": 9616 + }, + "sample_count": 9616, + "synchronized_e2e_ms": { + "max": 57.067035, + "median": 0.06512, + "min": 0.050337, + "p90": 0.1358085, + "sample_count": 9616 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.074016, + "submission_ms": 0.074016, + "synchronized_e2e_ms": 0.094272 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.047424, + "submission_ms": 0.047424, + "synchronized_e2e_ms": 0.063936 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.241056, + "submission_ms": 0.241056, + "synchronized_e2e_ms": 0.264064 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.947785, + "submission_ms": 8.947785, + "synchronized_e2e_ms": 8.987498 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.12304, + "submission_ms": 0.12304, + "synchronized_e2e_ms": 0.144448 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.044384, + "submission_ms": 0.044384, + "synchronized_e2e_ms": 0.062368 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.304322, + "submission_ms": 1.304322, + "synchronized_e2e_ms": 1.33165 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.165761, + "submission_ms": 1.165761, + "synchronized_e2e_ms": 1.189921 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.060224, + "submission_ms": 0.060224, + "synchronized_e2e_ms": 0.082336 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 41.918091, + "submission_ms": 41.918091, + "synchronized_e2e_ms": 41.951979 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.509327, + "evolution_kernel_ms": 0.169792, + "evolution_speedup": 2.9997, + "evolution_tflops": 0.1359, + "expected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d176_b1_n256_k256_d176", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 41716, + "measurement_schedule_sha256": "496ccaa32f787e03df88d2d3a4ca7e5fe0977709e2339088782021309738f5aa", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 11242, + "public_pair_count": 9616, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 11242, + "baseline_public_raw": 9616, + "candidate_precomputed": 11242, + "candidate_public_raw": 9616 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 8346 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:post_d895_d176_b1_n256_k256_d176", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.6177777777777778, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.795585073710074, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.21423299244119093, + "including_init_synchronized_e2e_speedup": 5.8459860898682905, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.2497982923950526, + "including_init_synchronized_e2e_speedup": 5.82363149792203, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.5580771955801686, + "including_init_synchronized_e2e_speedup": 5.61687572716988, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.7835674393545184, + "including_init_synchronized_e2e_speedup": 4.472128002480601, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 4.9797101449275365, + "hot_synchronized_e2e_speedup": 2.795585073710074, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 217601, + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d176_b1_n256_k256_d176", + "source": "tailpad_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 176, + "K": 8192, + "N": 512, + "baseline_07cf_adapter_bench_iters": 785, + "baseline_07cf_adapter_gpu_span_ms": 0.158176, + "baseline_07cf_adapter_host_enqueue_ms": 0.143424, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.031264, + "baseline_07cf_adapter_kernel_sum_ms": 0.126881, + "baseline_07cf_adapter_submission_ms": 0.143424, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.256128, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.126881 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.19056, + "synchronized_e2e_ms": 0.305568 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.158176 + }, + "host_enqueue_ms": { + "median": 0.143424 + }, + "inter_kernel_gap_ms": { + "median": 0.031264 + }, + "kernel_sum_ms": { + "median": 0.126881 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 785, + "submission_ms": { + "median": 0.143424 + }, + "synchronized_e2e_ms": { + "median": 0.256128 + } + }, + "baseline_07cf_precomputed_bench_iters": 775, + "baseline_07cf_precomputed_gpu_span_ms": 0.134177, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042432, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.134177, + "baseline_07cf_precomputed_submission_ms": 0.042432, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.183329, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.134177 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.049024, + "synchronized_e2e_ms": 0.161984 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.134177 + }, + "host_enqueue_ms": { + "median": 0.042432 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.134177 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 775, + "submission_ms": { + "median": 0.042432 + }, + "synchronized_e2e_ms": { + "median": 0.183329 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.1788607585502733, + "submission": 3.3800904977375565, + "synchronized_e2e": 1.397094840423501 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.857896, + "after_init_synchronized_e2e_ms_per_call": 7.946984, + "including_init_host_enqueue_ms_per_call": 42.085803, + "including_init_synchronized_e2e_ms_per_call": 42.292427, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9148712, + "after_init_synchronized_e2e_ms_per_call": 1.0252136, + "including_init_host_enqueue_ms_per_call": 4.3376619, + "including_init_synchronized_e2e_ms_per_call": 4.4597579000000005, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22056872, + "after_init_synchronized_e2e_ms_per_call": 0.33303656000000004, + "including_init_host_enqueue_ms_per_call": 0.56284779, + "including_init_synchronized_e2e_ms_per_call": 0.67649099, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.151138472, + "after_init_synchronized_e2e_ms_per_call": 0.26381885600000005, + "including_init_host_enqueue_ms_per_call": 0.185366379, + "including_init_synchronized_e2e_ms_per_call": 0.298164299, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.857896, + "synchronized_e2e_ms": 7.946984, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.396832, + "median": 0.158176, + "min": 0.150048, + "p90": 0.1629952, + "sample_count": 785 + }, + "host_enqueue_ms": { + "max": 0.412288, + "median": 0.143424, + "min": 0.126048, + "p90": 0.16364859999999998, + "sample_count": 785 + }, + "sample_count": 785, + "synchronized_e2e_ms": { + "max": 0.513472, + "median": 0.256128, + "min": 0.24064, + "p90": 0.2753152, + "sample_count": 785 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 775, + "candidate_precomputed_gpu_span_ms": 0.127232, + "candidate_precomputed_host_enqueue_ms": 0.056256, + "candidate_precomputed_inter_kernel_gap_ms": 0.004896, + "candidate_precomputed_kernel_sum_ms": 0.122273, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.056256, + "candidate_precomputed_synchronized_e2e_ms": 0.163521, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.122273 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.043904, + "synchronized_e2e_ms": 0.148832 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.127232 + }, + "host_enqueue_ms": { + "median": 0.056256 + }, + "inter_kernel_gap_ms": { + "median": 0.004896 + }, + "kernel_sum_ms": { + "median": 0.122273 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 775, + "submission_ms": { + "median": 0.056256 + }, + "synchronized_e2e_ms": { + "median": 0.163521 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb9e4d4710", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb9e4d4590" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.0042756539235411, + "submission": 0.8020477815699659, + "synchronized_e2e": 1.072204793268143 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 785, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.127776, + "candidate_public_raw_host_enqueue_ms": 0.04512, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.127936, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.04512, + "candidate_public_raw_synchronized_e2e_ms": 0.175328, + "candidate_public_raw_tflops_from_gpu_span": 11.554556473829201, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.127776 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.059936, + "synchronized_e2e_ms": 0.177184 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.127776 + }, + "host_enqueue_ms": { + "median": 0.04512 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.127936 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 785, + "submission_ms": { + "median": 0.04512 + }, + "synchronized_e2e_ms": { + "median": 0.175328 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.368355, + "after_init_synchronized_e2e_ms_per_call": 2.478403, + "including_init_host_enqueue_ms_per_call": 36.945895, + "including_init_synchronized_e2e_ms_per_call": 450.566195, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.27744350000000007, + "after_init_synchronized_e2e_ms_per_call": 0.4056355, + "including_init_host_enqueue_ms_per_call": 3.7351975000000004, + "including_init_synchronized_e2e_ms_per_call": 45.2144147, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06835235, + "after_init_synchronized_e2e_ms_per_call": 0.19835875000000003, + "including_init_host_enqueue_ms_per_call": 0.41412774999999996, + "including_init_synchronized_e2e_ms_per_call": 4.67923667, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.047443235, + "after_init_synchronized_e2e_ms_per_call": 0.177631075, + "including_init_host_enqueue_ms_per_call": 0.082020775, + "including_init_synchronized_e2e_ms_per_call": 0.625718867, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.368355, + "synchronized_e2e_ms": 2.478403, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.128416, + "median": 0.127776, + "min": 0.127104, + "p90": 0.128, + "sample_count": 785 + }, + "host_enqueue_ms": { + "max": 0.07856, + "median": 0.04512, + "min": 0.03664, + "p90": 0.049728, + "sample_count": 785 + }, + "sample_count": 785, + "synchronized_e2e_ms": { + "max": 0.215744, + "median": 0.175328, + "min": 0.167648, + "p90": 0.17934819999999999, + "sample_count": 785 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.063104, + "submission_ms": 0.063104, + "synchronized_e2e_ms": 0.175648 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.049024, + "submission_ms": 0.049024, + "synchronized_e2e_ms": 0.161984 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.19056, + "submission_ms": 0.19056, + "synchronized_e2e_ms": 0.305568 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.857896, + "submission_ms": 7.857896, + "synchronized_e2e_ms": 7.946984 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.070816, + "submission_ms": 0.070816, + "synchronized_e2e_ms": 0.166848 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.043904, + "submission_ms": 0.043904, + "synchronized_e2e_ms": 0.148832 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.229601, + "submission_ms": 1.229601, + "synchronized_e2e_ms": 1.254178 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.065729, + "submission_ms": 1.065729, + "synchronized_e2e_ms": 1.087841 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.059936, + "submission_ms": 0.059936, + "synchronized_e2e_ms": 0.177184 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.368355, + "submission_ms": 2.368355, + "synchronized_e2e_ms": 2.478403 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 2.050605, + "evolution_kernel_ms": 0.301759, + "evolution_speedup": 6.7955, + "evolution_tflops": 4.8926, + "expected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d176_b1_n512_k8192_d176", + "measurement_order": [ + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 3120, + "measurement_schedule_sha256": "15cbab796debb68bbb58b5264306086fa44bd56246bef60060db6c4be58ebf5d", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 775, + "public_pair_count": 785, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 775, + "baseline_public_raw": 785, + "candidate_precomputed": 775, + "candidate_public_raw": 785 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 624 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:post_d895_d176_b1_n512_k8192_d176", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.0545853244466799, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.4608505201679138, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.206493859150428, + "including_init_synchronized_e2e_speedup": 0.09386506903830191, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.527425730736092, + "including_init_synchronized_e2e_speedup": 0.09863575431841211, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 1.6789607718338615, + "including_init_synchronized_e2e_speedup": 0.14457293736330717, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.485206662178901, + "including_init_synchronized_e2e_speedup": 0.4765147971796734, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.2379163536188331, + "hot_synchronized_e2e_speedup": 1.4608505201679138, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 217604, + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d176_b1_n512_k8192_d176", + "source": "high_k_low_n", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 128, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 8 + } + }, + { + "B": 2, + "D": 176, + "K": 1024, + "N": 2048, + "baseline_07cf_adapter_bench_iters": 3377, + "baseline_07cf_adapter_gpu_span_ms": 0.068608, + "baseline_07cf_adapter_host_enqueue_ms": 0.146401, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.035104, + "baseline_07cf_adapter_kernel_sum_ms": 0.033472, + "baseline_07cf_adapter_submission_ms": 0.146401, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.169441, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.033472 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.27472, + "synchronized_e2e_ms": 0.298592 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.068608 + }, + "host_enqueue_ms": { + "median": 0.146401 + }, + "inter_kernel_gap_ms": { + "median": 0.035104 + }, + "kernel_sum_ms": { + "median": 0.033472 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3377, + "submission_ms": { + "median": 0.146401 + }, + "synchronized_e2e_ms": { + "median": 0.169441 + } + }, + "baseline_07cf_precomputed_bench_iters": 3643, + "baseline_07cf_precomputed_gpu_span_ms": 0.027456, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042976, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.027456, + "baseline_07cf_precomputed_submission_ms": 0.042976, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.077824, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.027456 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.04784, + "synchronized_e2e_ms": 0.0768 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.027456 + }, + "host_enqueue_ms": { + "median": 0.042976 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.027456 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3643, + "submission_ms": { + "median": 0.042976 + }, + "synchronized_e2e_ms": { + "median": 0.077824 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.4988344988344986, + "submission": 3.4065757632166793, + "synchronized_e2e": 2.1772332442434212 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.524873, + "after_init_synchronized_e2e_ms_per_call": 8.551913, + "including_init_host_enqueue_ms_per_call": 42.974317, + "including_init_synchronized_e2e_ms_per_call": 492.593407, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9842482, + "after_init_synchronized_e2e_ms_per_call": 1.0076882, + "including_init_host_enqueue_ms_per_call": 4.4291925999999995, + "including_init_synchronized_e2e_ms_per_call": 49.4118376, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23018571999999998, + "after_init_synchronized_e2e_ms_per_call": 0.25326571999999997, + "including_init_host_enqueue_ms_per_call": 0.57468016, + "including_init_synchronized_e2e_ms_per_call": 5.09368066, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.154779472, + "after_init_synchronized_e2e_ms_per_call": 0.177823472, + "including_init_host_enqueue_ms_per_call": 0.18922891600000002, + "including_init_synchronized_e2e_ms_per_call": 0.661864966, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.524873, + "synchronized_e2e_ms": 8.551913, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.252197, + "median": 0.068608, + "min": 0.05968, + "p90": 0.09063680000000002, + "sample_count": 3377 + }, + "host_enqueue_ms": { + "max": 99.095814, + "median": 0.146401, + "min": 0.121888, + "p90": 0.2134724, + "sample_count": 3377 + }, + "sample_count": 3377, + "synchronized_e2e_ms": { + "max": 99.324582, + "median": 0.169441, + "min": 0.146944, + "p90": 0.23564800000000002, + "sample_count": 3377 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3643, + "candidate_precomputed_gpu_span_ms": 0.02624, + "candidate_precomputed_host_enqueue_ms": 0.056128, + "candidate_precomputed_inter_kernel_gap_ms": 0.00608, + "candidate_precomputed_kernel_sum_ms": 0.020128, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.056128, + "candidate_precomputed_synchronized_e2e_ms": 0.06864, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.020128 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.051328, + "synchronized_e2e_ms": 0.066368 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.02624 + }, + "host_enqueue_ms": { + "median": 0.056128 + }, + "inter_kernel_gap_ms": { + "median": 0.00608 + }, + "kernel_sum_ms": { + "median": 0.020128 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3643, + "submission_ms": { + "median": 0.056128 + }, + "synchronized_e2e_ms": { + "median": 0.06864 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb724eb020", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb724ea7e0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.128048780487805, + "submission": 0.8112884834663626, + "synchronized_e2e": 1.1379953379953378 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 3377, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.0296, + "candidate_public_raw_host_enqueue_ms": 0.045536, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.029824, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.045536, + "candidate_public_raw_synchronized_e2e_ms": 0.078112, + "candidate_public_raw_tflops_from_gpu_span": 49.878209729729726, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.0296 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.062464, + "synchronized_e2e_ms": 0.086144 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.0296 + }, + "host_enqueue_ms": { + "median": 0.045536 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.029824 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3377, + "submission_ms": { + "median": 0.045536 + }, + "synchronized_e2e_ms": { + "median": 0.078112 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 48.374226, + "after_init_synchronized_e2e_ms_per_call": 48.416946, + "including_init_host_enqueue_ms_per_call": 83.24831, + "including_init_synchronized_e2e_ms_per_call": 83.360022, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 4.878405, + "after_init_synchronized_e2e_ms_per_call": 4.9119954, + "including_init_host_enqueue_ms_per_call": 8.3658134, + "including_init_synchronized_e2e_ms_per_call": 8.406303, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.5288229, + "after_init_synchronized_e2e_ms_per_call": 0.56150034, + "including_init_host_enqueue_ms_per_call": 0.87756374, + "including_init_synchronized_e2e_ms_per_call": 0.9109311, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.09386469, + "after_init_synchronized_e2e_ms_per_call": 0.126450834, + "including_init_host_enqueue_ms_per_call": 0.128738774, + "including_init_synchronized_e2e_ms_per_call": 0.16139391, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 48.374226, + "synchronized_e2e_ms": 48.416946, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.030464, + "median": 0.0296, + "min": 0.029184, + "p90": 0.030144, + "sample_count": 3377 + }, + "host_enqueue_ms": { + "max": 0.620577, + "median": 0.045536, + "min": 0.034528, + "p90": 0.0683652, + "sample_count": 3377 + }, + "sample_count": 3377, + "synchronized_e2e_ms": { + "max": 0.721057, + "median": 0.078112, + "min": 0.069536, + "p90": 0.09797760000000001, + "sample_count": 3377 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.07696, + "submission_ms": 0.07696, + "synchronized_e2e_ms": 0.10384 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.04784, + "submission_ms": 0.04784, + "synchronized_e2e_ms": 0.0768 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.27472, + "submission_ms": 0.27472, + "synchronized_e2e_ms": 0.298592 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.524873, + "submission_ms": 8.524873, + "synchronized_e2e_ms": 8.551913 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.091712, + "submission_ms": 0.091712, + "synchronized_e2e_ms": 0.111456 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.051328, + "submission_ms": 0.051328, + "synchronized_e2e_ms": 0.066368 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.596258, + "submission_ms": 1.596258, + "synchronized_e2e_ms": 1.626722 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.279394, + "submission_ms": 1.279394, + "synchronized_e2e_ms": 1.307106 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.062464, + "submission_ms": 0.062464, + "synchronized_e2e_ms": 0.086144 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 48.374226, + "submission_ms": 48.374226, + "synchronized_e2e_ms": 48.416946 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.486367, + "evolution_kernel_ms": 0.191839, + "evolution_speedup": 2.5353, + "evolution_tflops": 7.696, + "expected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d176_b2_n2048_k1024_d176", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 14040, + "measurement_schedule_sha256": "bdfb9f9d46dee8067d35f512a7c736c829e4e80e827f6b5de33720a82601e2ec", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3643, + "public_pair_count": 3377, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3643, + "baseline_public_raw": 3377, + "candidate_precomputed": 3643, + "candidate_public_raw": 3377 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2810 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:post_d895_d176_b2_n2048_k1024_d176", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.0463414634146342, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.1692057558377713, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.1766305747578544, + "including_init_synchronized_e2e_speedup": 5.909228370885027, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.20514844130350773, + "including_init_synchronized_e2e_speedup": 5.87795105648702, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.4510517660594826, + "including_init_synchronized_e2e_speedup": 5.591729890438476, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.406265711145883, + "including_init_synchronized_e2e_speedup": 4.1009290003569525, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.3178378378378377, + "hot_synchronized_e2e_speedup": 2.1692057558377713, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 217602, + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d176_b2_n2048_k1024_d176", + "source": "tailpad_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 176, + "K": 1024, + "N": 8192, + "baseline_07cf_adapter_bench_iters": 1507, + "baseline_07cf_adapter_gpu_span_ms": 0.106176, + "baseline_07cf_adapter_host_enqueue_ms": 0.173664, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.012128, + "baseline_07cf_adapter_kernel_sum_ms": 0.093728, + "baseline_07cf_adapter_submission_ms": 0.173664, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.222048, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.093728 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.210177, + "synchronized_e2e_ms": 0.278497 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.106176 + }, + "host_enqueue_ms": { + "median": 0.173664 + }, + "inter_kernel_gap_ms": { + "median": 0.012128 + }, + "kernel_sum_ms": { + "median": 0.093728 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1507, + "submission_ms": { + "median": 0.173664 + }, + "synchronized_e2e_ms": { + "median": 0.222048 + } + }, + "baseline_07cf_precomputed_bench_iters": 3014, + "baseline_07cf_precomputed_gpu_span_ms": 0.053888, + "baseline_07cf_precomputed_host_enqueue_ms": 0.051872, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.053888, + "baseline_07cf_precomputed_submission_ms": 0.051872, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.112736, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.053888 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.052704, + "synchronized_e2e_ms": 0.111585 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.053888 + }, + "host_enqueue_ms": { + "median": 0.051872 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.053888 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3014, + "submission_ms": { + "median": 0.051872 + }, + "synchronized_e2e_ms": { + "median": 0.112736 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.970308788598575, + "submission": 3.3479333744602098, + "synchronized_e2e": 1.9696281578200396 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.293799, + "after_init_synchronized_e2e_ms_per_call": 7.324999, + "including_init_host_enqueue_ms_per_call": 42.845067, + "including_init_synchronized_e2e_ms_per_call": 42.956236, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8856775000000001, + "after_init_synchronized_e2e_ms_per_call": 0.9323431, + "including_init_host_enqueue_ms_per_call": 4.4408043, + "including_init_synchronized_e2e_ms_per_call": 4.4954668, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.24486535, + "after_init_synchronized_e2e_ms_per_call": 0.29307751, + "including_init_host_enqueue_ms_per_call": 0.60037803, + "including_init_synchronized_e2e_ms_per_call": 0.64938988, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.180784135, + "after_init_synchronized_e2e_ms_per_call": 0.229150951, + "including_init_host_enqueue_ms_per_call": 0.216335403, + "including_init_synchronized_e2e_ms_per_call": 0.26478218800000003, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.293799, + "synchronized_e2e_ms": 7.324999, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.399136, + "median": 0.106176, + "min": 0.094304, + "p90": 0.12477440000000001, + "sample_count": 1507 + }, + "host_enqueue_ms": { + "max": 1.024225, + "median": 0.173664, + "min": 0.147584, + "p90": 0.2070848, + "sample_count": 1507 + }, + "sample_count": 1507, + "synchronized_e2e_ms": { + "max": 1.103361, + "median": 0.222048, + "min": 0.19808, + "p90": 0.25259580000000004, + "sample_count": 1507 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3014, + "candidate_precomputed_gpu_span_ms": 0.028896, + "candidate_precomputed_host_enqueue_ms": 0.063232, + "candidate_precomputed_inter_kernel_gap_ms": 0.003648, + "candidate_precomputed_kernel_sum_ms": 0.025248, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.063232, + "candidate_precomputed_synchronized_e2e_ms": 0.076384, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.025248 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.053408, + "synchronized_e2e_ms": 0.06976 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.028896 + }, + "host_enqueue_ms": { + "median": 0.063232 + }, + "inter_kernel_gap_ms": { + "median": 0.003648 + }, + "kernel_sum_ms": { + "median": 0.025248 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3014, + "submission_ms": { + "median": 0.063232 + }, + "synchronized_e2e_ms": { + "median": 0.076384 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb9111b530", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb9111a690" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.295681063122924, + "submission": 0.945344129554656, + "synchronized_e2e": 1.6790950984499373 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 1507, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.066336, + "candidate_public_raw_host_enqueue_ms": 0.059776, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.066272, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.059776, + "candidate_public_raw_synchronized_e2e_ms": 0.128256, + "candidate_public_raw_tflops_from_gpu_span": 178.05053159671974, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.066272 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.077568, + "synchronized_e2e_ms": 0.127008 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.066336 + }, + "host_enqueue_ms": { + "median": 0.059776 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.066272 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1507, + "submission_ms": { + "median": 0.059776 + }, + "synchronized_e2e_ms": { + "median": 0.128256 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.519811, + "after_init_synchronized_e2e_ms_per_call": 2.560963, + "including_init_host_enqueue_ms_per_call": 38.357831, + "including_init_synchronized_e2e_ms_per_call": 486.661176, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.3057795, + "after_init_synchronized_e2e_ms_per_call": 0.3715267, + "including_init_host_enqueue_ms_per_call": 3.8895815, + "including_init_synchronized_e2e_ms_per_call": 48.781548, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.08437635, + "after_init_synchronized_e2e_ms_per_call": 0.15258307000000002, + "including_init_host_enqueue_ms_per_call": 0.44275655, + "including_init_synchronized_e2e_ms_per_call": 4.9935852, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.062236035, + "after_init_synchronized_e2e_ms_per_call": 0.130688707, + "including_init_host_enqueue_ms_per_call": 0.09807405500000001, + "including_init_synchronized_e2e_ms_per_call": 0.61478892, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.519811, + "synchronized_e2e_ms": 2.560963, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.06736, + "median": 0.066336, + "min": 0.0656, + "p90": 0.066656, + "sample_count": 1507 + }, + "host_enqueue_ms": { + "max": 70.617737, + "median": 0.059776, + "min": 0.049856, + "p90": 0.0750208, + "sample_count": 1507 + }, + "sample_count": 1507, + "synchronized_e2e_ms": { + "max": 70.886697, + "median": 0.128256, + "min": 0.117536, + "p90": 0.14169640000000003, + "sample_count": 1507 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.06992, + "submission_ms": 0.06992, + "synchronized_e2e_ms": 0.126112 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.052704, + "submission_ms": 0.052704, + "synchronized_e2e_ms": 0.111585 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.210177, + "submission_ms": 0.210177, + "synchronized_e2e_ms": 0.278497 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.293799, + "submission_ms": 7.293799, + "synchronized_e2e_ms": 7.324999 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.087104, + "submission_ms": 0.087104, + "synchronized_e2e_ms": 0.104768 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.053408, + "submission_ms": 0.053408, + "synchronized_e2e_ms": 0.06976 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.229601, + "submission_ms": 1.229601, + "synchronized_e2e_ms": 1.254401 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.068865, + "submission_ms": 1.068865, + "synchronized_e2e_ms": 1.094049 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.077568, + "submission_ms": 0.077568, + "synchronized_e2e_ms": 0.127008 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.519811, + "submission_ms": 2.519811, + "synchronized_e2e_ms": 2.560963 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.492175, + "evolution_kernel_ms": 0.195712, + "evolution_speedup": 2.5148, + "evolution_tflops": 60.3497, + "expected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d176_b4_n8192_k1024_d176", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 9042, + "measurement_schedule_sha256": "81ec2f4e073a7aaf9a7b5c8606d5ebb519764f7b3a338e09f27a5c53f335fc7f", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3014, + "public_pair_count": 1507, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3014, + "baseline_public_raw": 1507, + "candidate_precomputed": 3014, + "candidate_public_raw": 1507 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 1810 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:post_d895_d176_b4_n8192_k1024_d176", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.8648947951273531, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.7312874251497004, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.860251788096899, + "including_init_synchronized_e2e_speedup": 0.08826723420402863, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.509491511646404, + "including_init_synchronized_e2e_speedup": 0.09215506650178465, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 1.920773451471385, + "including_init_synchronized_e2e_speedup": 0.13004481830008627, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.7534104993478894, + "including_init_synchronized_e2e_speedup": 0.43068796360220685, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.6005788712011577, + "hot_synchronized_e2e_speedup": 1.7312874251497004, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 217603, + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d176_b4_n8192_k1024_d176", + "source": "tailpad_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 224, + "K": 256, + "N": 256, + "baseline_07cf_adapter_bench_iters": 8866, + "baseline_07cf_adapter_gpu_span_ms": 0.063104, + "baseline_07cf_adapter_host_enqueue_ms": 0.179488, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.05232, + "baseline_07cf_adapter_kernel_sum_ms": 0.010784, + "baseline_07cf_adapter_submission_ms": 0.179488, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.20176, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.010784 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.176385, + "synchronized_e2e_ms": 0.197185 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.063104 + }, + "host_enqueue_ms": { + "median": 0.179488 + }, + "inter_kernel_gap_ms": { + "median": 0.05232 + }, + "kernel_sum_ms": { + "median": 0.010784 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 8866, + "submission_ms": { + "median": 0.179488 + }, + "synchronized_e2e_ms": { + "median": 0.20176 + } + }, + "baseline_07cf_precomputed_bench_iters": 9736, + "baseline_07cf_precomputed_gpu_span_ms": 0.010304, + "baseline_07cf_precomputed_host_enqueue_ms": 0.058464, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.010304, + "baseline_07cf_precomputed_submission_ms": 0.058464, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.074976, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.010304 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.040448, + "synchronized_e2e_ms": 0.056224 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.010304 + }, + "host_enqueue_ms": { + "median": 0.058464 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.010304 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 9736, + "submission_ms": { + "median": 0.058464 + }, + "synchronized_e2e_ms": { + "median": 0.074976 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 6.124223602484471, + "submission": 3.0700602079912427, + "synchronized_e2e": 2.690994451557832 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 6.660263, + "after_init_synchronized_e2e_ms_per_call": 6.688455, + "including_init_host_enqueue_ms_per_call": 43.594701, + "including_init_synchronized_e2e_ms_per_call": 461.799774, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8275655000000001, + "after_init_synchronized_e2e_ms_per_call": 0.8504295000000001, + "including_init_host_enqueue_ms_per_call": 4.5210093, + "including_init_synchronized_e2e_ms_per_call": 46.3615614, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.24429575, + "after_init_synchronized_e2e_ms_per_call": 0.26662695, + "including_init_host_enqueue_ms_per_call": 0.61364013, + "including_init_synchronized_e2e_ms_per_call": 4.817740140000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.185968775, + "after_init_synchronized_e2e_ms_per_call": 0.20824669499999998, + "including_init_host_enqueue_ms_per_call": 0.222903213, + "including_init_synchronized_e2e_ms_per_call": 0.663358014, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 6.660263, + "synchronized_e2e_ms": 6.688455, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.600293, + "median": 0.063104, + "min": 0.044448, + "p90": 0.097248, + "sample_count": 8866 + }, + "host_enqueue_ms": { + "max": 80.451059, + "median": 0.179488, + "min": 0.12464, + "p90": 0.386432, + "sample_count": 8866 + }, + "sample_count": 8866, + "synchronized_e2e_ms": { + "max": 85.635705, + "median": 0.20176, + "min": 0.141952, + "p90": 0.4401605, + "sample_count": 8866 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 9736, + "candidate_precomputed_gpu_span_ms": 0.015648, + "candidate_precomputed_host_enqueue_ms": 0.064864, + "candidate_precomputed_inter_kernel_gap_ms": 0.006912, + "candidate_precomputed_kernel_sum_ms": 0.008768, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.064864, + "candidate_precomputed_synchronized_e2e_ms": 0.079616, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.008768 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.045536, + "synchronized_e2e_ms": 0.061344 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.015648 + }, + "host_enqueue_ms": { + "median": 0.064864 + }, + "inter_kernel_gap_ms": { + "median": 0.006912 + }, + "kernel_sum_ms": { + "median": 0.008768 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 9736, + "submission_ms": { + "median": 0.064864 + }, + "synchronized_e2e_ms": { + "median": 0.079616 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01af30e0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01af3650" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 0.736196319018405, + "submission": 0.8855451406018746, + "synchronized_e2e": 0.929065765675241 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 8866, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.01152, + "candidate_public_raw_host_enqueue_ms": 0.05744, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.011712, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.05744, + "candidate_public_raw_synchronized_e2e_ms": 0.07396849999999999, + "candidate_public_raw_tflops_from_gpu_span": 2.548622222222222, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.01152 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.078016, + "synchronized_e2e_ms": 0.097888 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.01152 + }, + "host_enqueue_ms": { + "median": 0.05744 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.011712 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 8866, + "submission_ms": { + "median": 0.05744 + }, + "synchronized_e2e_ms": { + "median": 0.07396849999999999 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 39.16692, + "after_init_synchronized_e2e_ms_per_call": 39.198664, + "including_init_host_enqueue_ms_per_call": 76.524943, + "including_init_synchronized_e2e_ms_per_call": 76.634287, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 3.9683879999999996, + "after_init_synchronized_e2e_ms_per_call": 3.9864380500000003, + "including_init_host_enqueue_ms_per_call": 7.704190299999999, + "including_init_synchronized_e2e_ms_per_call": 7.73000035, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.44853479999999996, + "after_init_synchronized_e2e_ms_per_call": 0.465215455, + "including_init_host_enqueue_ms_per_call": 0.82211503, + "including_init_synchronized_e2e_ms_per_call": 0.8395716849999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.09654948, + "after_init_synchronized_e2e_ms_per_call": 0.11309319550000001, + "including_init_host_enqueue_ms_per_call": 0.13390750299999998, + "including_init_synchronized_e2e_ms_per_call": 0.1505288185, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 39.16692, + "synchronized_e2e_ms": 39.198664, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.012224, + "median": 0.01152, + "min": 0.01104, + "p90": 0.011808, + "sample_count": 8866 + }, + "host_enqueue_ms": { + "max": 51.647413, + "median": 0.05744, + "min": 0.035712, + "p90": 0.15384, + "sample_count": 8866 + }, + "sample_count": 8866, + "synchronized_e2e_ms": { + "max": 51.851221, + "median": 0.07396849999999999, + "min": 0.048672, + "p90": 0.20316800000000002, + "sample_count": 8866 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.06448, + "submission_ms": 0.06448, + "synchronized_e2e_ms": 0.081152 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.040448, + "submission_ms": 0.040448, + "synchronized_e2e_ms": 0.056224 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.176385, + "submission_ms": 0.176385, + "synchronized_e2e_ms": 0.197185 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 6.660263, + "submission_ms": 6.660263, + "synchronized_e2e_ms": 6.688455 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.094176, + "submission_ms": 0.094176, + "synchronized_e2e_ms": 0.113408 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.045536, + "submission_ms": 0.045536, + "synchronized_e2e_ms": 0.061344 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.257441, + "submission_ms": 1.257441, + "synchronized_e2e_ms": 1.283521 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.108257, + "submission_ms": 1.108257, + "synchronized_e2e_ms": 1.129537 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.078016, + "submission_ms": 0.078016, + "synchronized_e2e_ms": 0.097888 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 39.16692, + "submission_ms": 39.16692, + "synchronized_e2e_ms": 39.198664 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.480639, + "evolution_kernel_ms": 0.170496, + "evolution_speedup": 2.8191, + "evolution_tflops": 0.1722, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d224_b1_n256_k256_d224", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 37204, + "measurement_schedule_sha256": "a7bae53515356827773da67cdf4dcabe0b9cefc831d8e16c8e2074f26edba543", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 9736, + "public_pair_count": 8866, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 9736, + "baseline_public_raw": 8866, + "candidate_precomputed": 9736, + "candidate_public_raw": 8866 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 7444 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:post_d895_d224_b1_n256_k256_d224", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.6584867075664622, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.727647579712986, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.17062966737845964, + "including_init_synchronized_e2e_speedup": 6.0260203634438465, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.21333066996989958, + "including_init_synchronized_e2e_speedup": 5.997614398555622, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.5731257359023035, + "including_init_synchronized_e2e_speedup": 5.7383308966642925, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.8413724546318966, + "including_init_synchronized_e2e_speedup": 4.406850599176131, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 5.477777777777777, + "hot_synchronized_e2e_speedup": 2.727647579712986, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 222404, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d224_b1_n256_k256_d224", + "source": "medium_small_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 224, + "K": 8192, + "N": 512, + "baseline_07cf_adapter_bench_iters": 1195, + "baseline_07cf_adapter_gpu_span_ms": 0.156961, + "baseline_07cf_adapter_host_enqueue_ms": 0.148512, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.032096, + "baseline_07cf_adapter_kernel_sum_ms": 0.124896, + "baseline_07cf_adapter_submission_ms": 0.148512, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.259008, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.124896 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.210241, + "synchronized_e2e_ms": 0.312225 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.156961 + }, + "host_enqueue_ms": { + "median": 0.148512 + }, + "inter_kernel_gap_ms": { + "median": 0.032096 + }, + "kernel_sum_ms": { + "median": 0.124896 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1195, + "submission_ms": { + "median": 0.148512 + }, + "synchronized_e2e_ms": { + "median": 0.259008 + } + }, + "baseline_07cf_precomputed_bench_iters": 1283, + "baseline_07cf_precomputed_gpu_span_ms": 0.134496, + "baseline_07cf_precomputed_host_enqueue_ms": 0.04368, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.134496, + "baseline_07cf_precomputed_submission_ms": 0.04368, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.184864, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.134496 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.043168, + "synchronized_e2e_ms": 0.150208 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.134496 + }, + "host_enqueue_ms": { + "median": 0.04368 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.134496 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1283, + "submission_ms": { + "median": 0.04368 + }, + "synchronized_e2e_ms": { + "median": 0.184864 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.1670309897692124, + "submission": 3.4000000000000004, + "synchronized_e2e": 1.401073221395188 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.334633, + "after_init_synchronized_e2e_ms_per_call": 8.402025, + "including_init_host_enqueue_ms_per_call": 42.56254, + "including_init_synchronized_e2e_ms_per_call": 42.747468000000005, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9671241, + "after_init_synchronized_e2e_ms_per_call": 1.0733097, + "including_init_host_enqueue_ms_per_call": 4.3899148, + "including_init_synchronized_e2e_ms_per_call": 4.507854, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23037321, + "after_init_synchronized_e2e_ms_per_call": 0.34043817000000004, + "including_init_host_enqueue_ms_per_call": 0.57265228, + "including_init_synchronized_e2e_ms_per_call": 0.6838926000000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15669812100000002, + "after_init_synchronized_e2e_ms_per_call": 0.26715101700000005, + "including_init_host_enqueue_ms_per_call": 0.19092602800000003, + "including_init_synchronized_e2e_ms_per_call": 0.3014964600000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.334633, + "synchronized_e2e_ms": 8.402025, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.534369, + "median": 0.156961, + "min": 0.149152, + "p90": 0.1734082, + "sample_count": 1195 + }, + "host_enqueue_ms": { + "max": 0.538272, + "median": 0.148512, + "min": 0.127296, + "p90": 0.18679680000000004, + "sample_count": 1195 + }, + "sample_count": 1195, + "synchronized_e2e_ms": { + "max": 0.634944, + "median": 0.259008, + "min": 0.239712, + "p90": 0.29407360000000005, + "sample_count": 1195 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 1283, + "candidate_precomputed_gpu_span_ms": 0.075936, + "candidate_precomputed_host_enqueue_ms": 0.059008, + "candidate_precomputed_inter_kernel_gap_ms": 0.004384, + "candidate_precomputed_kernel_sum_ms": 0.071489, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.059008, + "candidate_precomputed_synchronized_e2e_ms": 0.113888, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.071489 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.0464, + "synchronized_e2e_ms": 0.10832 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.075936 + }, + "host_enqueue_ms": { + "median": 0.059008 + }, + "inter_kernel_gap_ms": { + "median": 0.004384 + }, + "kernel_sum_ms": { + "median": 0.071489 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1283, + "submission_ms": { + "median": 0.059008 + }, + "synchronized_e2e_ms": { + "median": 0.113888 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc00a39610", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc00a3a990" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.103244837758112, + "submission": 0.7868763557483731, + "synchronized_e2e": 1.158752458555774 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 1195, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.083776, + "candidate_public_raw_host_enqueue_ms": 0.046432, + "candidate_public_raw_inter_kernel_gap_ms": 0.000128, + "candidate_public_raw_kernel_sum_ms": 0.083649, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.046432, + "candidate_public_raw_synchronized_e2e_ms": 0.131968, + "candidate_public_raw_tflops_from_gpu_span": 22.42943315508021, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.083649 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.057248, + "synchronized_e2e_ms": 0.132288 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.083776 + }, + "host_enqueue_ms": { + "median": 0.046432 + }, + "inter_kernel_gap_ms": { + "median": 0.000128 + }, + "kernel_sum_ms": { + "median": 0.083649 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1195, + "submission_ms": { + "median": 0.046432 + }, + "synchronized_e2e_ms": { + "median": 0.131968 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 38.945704, + "after_init_synchronized_e2e_ms_per_call": 39.010536, + "including_init_host_enqueue_ms_per_call": 73.523244, + "including_init_synchronized_e2e_ms_per_call": 487.098328, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 3.9363591999999996, + "after_init_synchronized_e2e_ms_per_call": 4.0198248, + "including_init_host_enqueue_ms_per_call": 7.394113200000001, + "including_init_synchronized_e2e_ms_per_call": 48.828604, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.43542472, + "after_init_synchronized_e2e_ms_per_call": 0.52075368, + "including_init_host_enqueue_ms_per_call": 0.78120012, + "including_init_synchronized_e2e_ms_per_call": 5.0016316, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.085331272, + "after_init_synchronized_e2e_ms_per_call": 0.170846568, + "including_init_host_enqueue_ms_per_call": 0.11990881200000002, + "including_init_synchronized_e2e_ms_per_call": 0.61893436, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 38.945704, + "synchronized_e2e_ms": 39.010536, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.08464, + "median": 0.083776, + "min": 0.083456, + "p90": 0.084064, + "sample_count": 1195 + }, + "host_enqueue_ms": { + "max": 0.119456, + "median": 0.046432, + "min": 0.03632, + "p90": 0.06314880000000003, + "sample_count": 1195 + }, + "sample_count": 1195, + "synchronized_e2e_ms": { + "max": 0.203264, + "median": 0.131968, + "min": 0.124224, + "p90": 0.14508200000000002, + "sample_count": 1195 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.064768, + "submission_ms": 0.064768, + "synchronized_e2e_ms": 0.17008 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.043168, + "submission_ms": 0.043168, + "synchronized_e2e_ms": 0.150208 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.210241, + "submission_ms": 0.210241, + "synchronized_e2e_ms": 0.312225 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.334633, + "submission_ms": 8.334633, + "synchronized_e2e_ms": 8.402025 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.071808, + "submission_ms": 0.071808, + "synchronized_e2e_ms": 0.127265 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.0464, + "submission_ms": 0.0464, + "synchronized_e2e_ms": 0.10832 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.283553, + "submission_ms": 1.283553, + "synchronized_e2e_ms": 1.308961 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.148577, + "submission_ms": 1.148577, + "synchronized_e2e_ms": 1.172385 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.057248, + "submission_ms": 0.057248, + "synchronized_e2e_ms": 0.132288 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 38.945704, + "submission_ms": 38.945704, + "synchronized_e2e_ms": 39.010536 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 2.069245, + "evolution_kernel_ms": 0.236528, + "evolution_speedup": 8.7484, + "evolution_tflops": 7.9443, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d224_b1_n512_k8192_d224", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 4956, + "measurement_schedule_sha256": "1d78dba4cd3ff078622e946fdc0d5f6dcc583f0f78e3446ef4ea6ce245fcfdc9", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 1283, + "public_pair_count": 1195, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 1283, + "baseline_public_raw": 1195, + "candidate_precomputed": 1283, + "candidate_public_raw": 1195 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 992 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:post_d895_d224_b1_n512_k8192_d224", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.7711757269279393, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.9626576139670224, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.215378353171051, + "including_init_synchronized_e2e_speedup": 0.08775942256981019, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.2670040992831329, + "including_init_synchronized_e2e_speedup": 0.09231994426873233, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.6537412659282601, + "including_init_synchronized_e2e_speedup": 0.13673390099342786, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.5636896902722683, + "including_init_synchronized_e2e_speedup": 0.48712186539457925, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.8735795454545452, + "hot_synchronized_e2e_speedup": 1.9626576139670224, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 222403, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d224_b1_n512_k8192_d224", + "source": "high_k_low_n", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 128, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 8 + } + }, + { + "B": 2, + "D": 224, + "K": 1024, + "N": 2048, + "baseline_07cf_adapter_bench_iters": 3979, + "baseline_07cf_adapter_gpu_span_ms": 0.070144, + "baseline_07cf_adapter_host_enqueue_ms": 0.14656, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.035264, + "baseline_07cf_adapter_kernel_sum_ms": 0.034848, + "baseline_07cf_adapter_submission_ms": 0.14656, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.171168, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.034848 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.187232, + "synchronized_e2e_ms": 0.208224 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.070144 + }, + "host_enqueue_ms": { + "median": 0.14656 + }, + "inter_kernel_gap_ms": { + "median": 0.035264 + }, + "kernel_sum_ms": { + "median": 0.034848 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3979, + "submission_ms": { + "median": 0.14656 + }, + "synchronized_e2e_ms": { + "median": 0.171168 + } + }, + "baseline_07cf_precomputed_bench_iters": 4134, + "baseline_07cf_precomputed_gpu_span_ms": 0.031008, + "baseline_07cf_precomputed_host_enqueue_ms": 0.0429765, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.031008, + "baseline_07cf_precomputed_submission_ms": 0.0429765, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.0813605, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.031008 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.04416, + "synchronized_e2e_ms": 0.072192 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.031008 + }, + "host_enqueue_ms": { + "median": 0.0429765 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.031008 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4134, + "submission_ms": { + "median": 0.0429765 + }, + "synchronized_e2e_ms": { + "median": 0.0813605 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.262125902992776, + "submission": 3.4102358265563737, + "synchronized_e2e": 2.103821879167409 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.701288, + "after_init_synchronized_e2e_ms_per_call": 7.728776, + "including_init_host_enqueue_ms_per_call": 42.150732, + "including_init_synchronized_e2e_ms_per_call": 491.77027, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9020328, + "after_init_synchronized_e2e_ms_per_call": 0.9269288, + "including_init_host_enqueue_ms_per_call": 4.3469771999999995, + "including_init_synchronized_e2e_ms_per_call": 49.33107819999999, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22210728, + "after_init_synchronized_e2e_ms_per_call": 0.24674408, + "including_init_host_enqueue_ms_per_call": 0.56660172, + "including_init_synchronized_e2e_ms_per_call": 5.08715902, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.154114728, + "after_init_synchronized_e2e_ms_per_call": 0.17872560799999998, + "including_init_host_enqueue_ms_per_call": 0.188564172, + "including_init_synchronized_e2e_ms_per_call": 0.662767102, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.701288, + "synchronized_e2e_ms": 7.728776, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.502305, + "median": 0.070144, + "min": 0.061056, + "p90": 0.0918144, + "sample_count": 3979 + }, + "host_enqueue_ms": { + "max": 16.1748, + "median": 0.14656, + "min": 0.12272, + "p90": 0.21645440000000002, + "sample_count": 3979 + }, + "sample_count": 3979, + "synchronized_e2e_ms": { + "max": 16.211024, + "median": 0.171168, + "min": 0.148288, + "p90": 0.2407552000000003, + "sample_count": 3979 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 4134, + "candidate_precomputed_gpu_span_ms": 0.02128, + "candidate_precomputed_host_enqueue_ms": 0.0564805, + "candidate_precomputed_inter_kernel_gap_ms": 0.005856, + "candidate_precomputed_kernel_sum_ms": 0.015424, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.0564805, + "candidate_precomputed_synchronized_e2e_ms": 0.068896, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.015424 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.044704, + "synchronized_e2e_ms": 0.061216 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.02128 + }, + "host_enqueue_ms": { + "median": 0.0564805 + }, + "inter_kernel_gap_ms": { + "median": 0.005856 + }, + "kernel_sum_ms": { + "median": 0.015424 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4134, + "submission_ms": { + "median": 0.0564805 + }, + "synchronized_e2e_ms": { + "median": 0.068896 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc0073ec60", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc0073e3f0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.1819548872180452, + "submission": 0.8045254556882464, + "synchronized_e2e": 1.0687412912215515 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 3979, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.025152, + "candidate_public_raw_host_enqueue_ms": 0.04544, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.025376, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.04544, + "candidate_public_raw_synchronized_e2e_ms": 0.073632, + "candidate_public_raw_tflops_from_gpu_span": 74.7077048346056, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.025152 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.07296, + "synchronized_e2e_ms": 0.091776 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.025152 + }, + "host_enqueue_ms": { + "median": 0.04544 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.025376 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 3979, + "submission_ms": { + "median": 0.04544 + }, + "synchronized_e2e_ms": { + "median": 0.073632 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 47.656722, + "after_init_synchronized_e2e_ms_per_call": 47.686514, + "including_init_host_enqueue_ms_per_call": 82.530806, + "including_init_synchronized_e2e_ms_per_call": 82.62959000000001, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 4.8065682, + "after_init_synchronized_e2e_ms_per_call": 4.834920200000001, + "including_init_host_enqueue_ms_per_call": 8.293976599999999, + "including_init_synchronized_e2e_ms_per_call": 8.329227800000002, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.52155282, + "after_init_synchronized_e2e_ms_per_call": 0.54976082, + "including_init_host_enqueue_ms_per_call": 0.87029366, + "including_init_synchronized_e2e_ms_per_call": 0.8991915800000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.093051282, + "after_init_synchronized_e2e_ms_per_call": 0.121244882, + "including_init_host_enqueue_ms_per_call": 0.12792536599999998, + "including_init_synchronized_e2e_ms_per_call": 0.15618795800000002, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 47.656722, + "synchronized_e2e_ms": 47.686514, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.02576, + "median": 0.025152, + "min": 0.024864, + "p90": 0.025376, + "sample_count": 3979 + }, + "host_enqueue_ms": { + "max": 15.533521, + "median": 0.04544, + "min": 0.035936, + "p90": 0.06960640000000001, + "sample_count": 3979 + }, + "sample_count": 3979, + "synchronized_e2e_ms": { + "max": 20.692246, + "median": 0.073632, + "min": 0.065312, + "p90": 0.09468819999999999, + "sample_count": 3979 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.066528, + "submission_ms": 0.066528, + "synchronized_e2e_ms": 0.093856 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.04416, + "submission_ms": 0.04416, + "synchronized_e2e_ms": 0.072192 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.187232, + "submission_ms": 0.187232, + "synchronized_e2e_ms": 0.208224 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.701288, + "submission_ms": 7.701288, + "synchronized_e2e_ms": 7.728776 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.087072, + "submission_ms": 0.087072, + "synchronized_e2e_ms": 0.10416 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.044704, + "submission_ms": 0.044704, + "synchronized_e2e_ms": 0.061216 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.486562, + "submission_ms": 1.486562, + "synchronized_e2e_ms": 1.509602 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.167873, + "submission_ms": 1.167873, + "synchronized_e2e_ms": 1.188897 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.07296, + "submission_ms": 0.07296, + "synchronized_e2e_ms": 0.091776 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 47.656722, + "submission_ms": 47.656722, + "synchronized_e2e_ms": 47.686514 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.47112, + "evolution_kernel_ms": 0.176704, + "evolution_speedup": 2.6662, + "evolution_tflops": 10.6339, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d224_b2_n2048_k1024_d224", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 16226, + "measurement_schedule_sha256": "2f1c809b82b68c180066edd9dc945376586964a9ae3c10049cb115a393b89284", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 4134, + "public_pair_count": 3979, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 4134, + "baseline_public_raw": 3979, + "candidate_precomputed": 4134, + "candidate_public_raw": 3979 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 3246 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:post_d895_d224_b2_n2048_k1024_d224", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.457142857142857, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.3246414602346803, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.16207466958058622, + "including_init_synchronized_e2e_speedup": 5.951503208475317, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.19171542893303592, + "including_init_synchronized_e2e_speedup": 5.92264725908925, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.448820779916619, + "including_init_synchronized_e2e_speedup": 5.657480711729973, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.4740878546939407, + "including_init_synchronized_e2e_speedup": 4.24339437231134, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.7888040712468194, + "hot_synchronized_e2e_speedup": 2.3246414602346803, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 222401, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d224_b2_n2048_k1024_d224", + "source": "new_medium_gap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 224, + "K": 1024, + "N": 8192, + "baseline_07cf_adapter_bench_iters": 1360, + "baseline_07cf_adapter_gpu_span_ms": 0.1074085, + "baseline_07cf_adapter_host_enqueue_ms": 0.170064, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.010048, + "baseline_07cf_adapter_kernel_sum_ms": 0.097312, + "baseline_07cf_adapter_submission_ms": 0.170064, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.2208485, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.097312 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.23568, + "synchronized_e2e_ms": 0.284032 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.1074085 + }, + "host_enqueue_ms": { + "median": 0.170064 + }, + "inter_kernel_gap_ms": { + "median": 0.010048 + }, + "kernel_sum_ms": { + "median": 0.097312 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1360, + "submission_ms": { + "median": 0.170064 + }, + "synchronized_e2e_ms": { + "median": 0.2208485 + } + }, + "baseline_07cf_precomputed_bench_iters": 2353, + "baseline_07cf_precomputed_gpu_span_ms": 0.057664, + "baseline_07cf_precomputed_host_enqueue_ms": 0.051904, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.057664, + "baseline_07cf_precomputed_submission_ms": 0.051904, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.116257, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.057664 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.0608, + "synchronized_e2e_ms": 0.123776 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.057664 + }, + "host_enqueue_ms": { + "median": 0.051904 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.057664 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2353, + "submission_ms": { + "median": 0.051904 + }, + "synchronized_e2e_ms": { + "median": 0.116257 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.8626612791342954, + "submission": 3.2765104808877927, + "synchronized_e2e": 1.8996576550229234 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.436072, + "after_init_synchronized_e2e_ms_per_call": 8.471016, + "including_init_host_enqueue_ms_per_call": 43.98734, + "including_init_synchronized_e2e_ms_per_call": 44.102253, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9966647999999999, + "after_init_synchronized_e2e_ms_per_call": 1.04586525, + "including_init_host_enqueue_ms_per_call": 4.5517916, + "including_init_synchronized_e2e_ms_per_call": 4.60898895, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.25272407999999996, + "after_init_synchronized_e2e_ms_per_call": 0.303350175, + "including_init_host_enqueue_ms_per_call": 0.60823676, + "including_init_synchronized_e2e_ms_per_call": 0.6596625449999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.17833000799999998, + "after_init_synchronized_e2e_ms_per_call": 0.2290986675, + "including_init_host_enqueue_ms_per_call": 0.213881276, + "including_init_synchronized_e2e_ms_per_call": 0.26472990450000006, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.436072, + "synchronized_e2e_ms": 8.471016, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 50.005908, + "median": 0.1074085, + "min": 0.097696, + "p90": 0.12876250000000003, + "sample_count": 1360 + }, + "host_enqueue_ms": { + "max": 50.586804, + "median": 0.170064, + "min": 0.144193, + "p90": 0.2071584000000001, + "sample_count": 1360 + }, + "sample_count": 1360, + "synchronized_e2e_ms": { + "max": 50.73954, + "median": 0.2208485, + "min": 0.198049, + "p90": 0.25512650000000003, + "sample_count": 1360 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 2353, + "candidate_precomputed_gpu_span_ms": 0.034112, + "candidate_precomputed_host_enqueue_ms": 0.062144, + "candidate_precomputed_inter_kernel_gap_ms": 0.00208, + "candidate_precomputed_kernel_sum_ms": 0.032032, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.062144, + "candidate_precomputed_synchronized_e2e_ms": 0.076448, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.032032 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.06416, + "synchronized_e2e_ms": 0.080832 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.034112 + }, + "host_enqueue_ms": { + "median": 0.062144 + }, + "inter_kernel_gap_ms": { + "median": 0.00208 + }, + "kernel_sum_ms": { + "median": 0.032032 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2353, + "submission_ms": { + "median": 0.062144 + }, + "synchronized_e2e_ms": { + "median": 0.076448 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb9f540440", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb9f5431d0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.1622889305816133, + "submission": 0.966529351184346, + "synchronized_e2e": 1.778777731268313 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 1360, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.07376, + "candidate_public_raw_host_enqueue_ms": 0.060064, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.073728, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.060064, + "candidate_public_raw_synchronized_e2e_ms": 0.135984, + "candidate_public_raw_tflops_from_gpu_span": 203.8013223427332, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.073728 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.126529, + "synchronized_e2e_ms": 0.180353 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.07376 + }, + "host_enqueue_ms": { + "median": 0.060064 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.073728 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1360, + "submission_ms": { + "median": 0.060064 + }, + "synchronized_e2e_ms": { + "median": 0.135984 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.212292, + "after_init_synchronized_e2e_ms_per_call": 3.258564, + "including_init_host_enqueue_ms_per_call": 39.050312, + "including_init_synchronized_e2e_ms_per_call": 487.358777, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.37528680000000003, + "after_init_synchronized_e2e_ms_per_call": 0.4482419999999999, + "including_init_host_enqueue_ms_per_call": 3.9590888, + "including_init_synchronized_e2e_ms_per_call": 48.8582633, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.09158628, + "after_init_synchronized_e2e_ms_per_call": 0.16720979999999996, + "including_init_host_enqueue_ms_per_call": 0.44996648, + "including_init_synchronized_e2e_ms_per_call": 5.00821193, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.063216228, + "after_init_synchronized_e2e_ms_per_call": 0.13910658, + "including_init_host_enqueue_ms_per_call": 0.099054248, + "including_init_synchronized_e2e_ms_per_call": 0.6232067929999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.212292, + "synchronized_e2e_ms": 3.258564, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.07472, + "median": 0.07376, + "min": 0.07312, + "p90": 0.07408, + "sample_count": 1360 + }, + "host_enqueue_ms": { + "max": 32.343074, + "median": 0.060064, + "min": 0.045376, + "p90": 0.07922560000000003, + "sample_count": 1360 + }, + "sample_count": 1360, + "synchronized_e2e_ms": { + "max": 32.407426, + "median": 0.135984, + "min": 0.12208, + "p90": 0.1523561, + "sample_count": 1360 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.083584, + "submission_ms": 0.083584, + "synchronized_e2e_ms": 0.144416 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.0608, + "submission_ms": 0.0608, + "synchronized_e2e_ms": 0.123776 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.23568, + "submission_ms": 0.23568, + "synchronized_e2e_ms": 0.284032 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.436072, + "submission_ms": 8.436072, + "synchronized_e2e_ms": 8.471016 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.127712, + "submission_ms": 0.127712, + "synchronized_e2e_ms": 0.151104 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.06416, + "submission_ms": 0.06416, + "synchronized_e2e_ms": 0.080832 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.332001, + "submission_ms": 1.332001, + "synchronized_e2e_ms": 1.357409 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.099394, + "submission_ms": 1.099394, + "synchronized_e2e_ms": 1.122818 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.126529, + "submission_ms": 0.126529, + "synchronized_e2e_ms": 0.180353 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.212292, + "submission_ms": 3.212292, + "synchronized_e2e_ms": 3.258564 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.553311, + "evolution_kernel_ms": 0.190783, + "evolution_speedup": 2.9002, + "evolution_tflops": 78.7929, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d224_b4_n8192_k1024_d224", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 7426, + "measurement_schedule_sha256": "63128f0b8683a75fc0880698158fe386b4d0b1df588f22a32755c23ef28088d2", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 2353, + "public_pair_count": 1360, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 2353, + "baseline_public_raw": 1360, + "candidate_precomputed": 2353, + "candidate_public_raw": 1360 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 1486 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:post_d895_d224_b4_n8192_k1024_d224", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.6904315196998123, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.6240770973055654, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.5996162726894427, + "including_init_synchronized_e2e_speedup": 0.090492374573568, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.3332602701219436, + "including_init_synchronized_e2e_speedup": 0.09433386777789951, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 1.8141889709813663, + "including_init_synchronized_e2e_speedup": 0.13171618018968298, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.6469290489349964, + "including_init_synchronized_e2e_speedup": 0.4247866157325408, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.4561889913232104, + "hot_synchronized_e2e_speedup": 1.6240770973055654, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 222402, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d224_b4_n8192_k1024_d224", + "source": "new_medium_gap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 256, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 288, + "K": 256, + "N": 256, + "baseline_07cf_adapter_bench_iters": 7071, + "baseline_07cf_adapter_gpu_span_ms": 0.053824, + "baseline_07cf_adapter_host_enqueue_ms": 0.145409, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.041344, + "baseline_07cf_adapter_kernel_sum_ms": 0.012448, + "baseline_07cf_adapter_submission_ms": 0.145409, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.165984, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.012448 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.198177, + "synchronized_e2e_ms": 0.216673 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.053824 + }, + "host_enqueue_ms": { + "median": 0.145409 + }, + "inter_kernel_gap_ms": { + "median": 0.041344 + }, + "kernel_sum_ms": { + "median": 0.012448 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 7071, + "submission_ms": { + "median": 0.145409 + }, + "synchronized_e2e_ms": { + "median": 0.165984 + } + }, + "baseline_07cf_precomputed_bench_iters": 8401, + "baseline_07cf_precomputed_gpu_span_ms": 0.011936, + "baseline_07cf_precomputed_host_enqueue_ms": 0.043232, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.011936, + "baseline_07cf_precomputed_submission_ms": 0.043232, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.060832, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.011936 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.038304, + "synchronized_e2e_ms": 0.052672 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.011936 + }, + "host_enqueue_ms": { + "median": 0.043232 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.011936 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 8401, + "submission_ms": { + "median": 0.043232 + }, + "synchronized_e2e_ms": { + "median": 0.060832 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 4.509383378016086, + "submission": 3.3634576239822356, + "synchronized_e2e": 2.728563913729616 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.923112, + "after_init_synchronized_e2e_ms_per_call": 7.949544, + "including_init_host_enqueue_ms_per_call": 44.85755, + "including_init_synchronized_e2e_ms_per_call": 463.06086300000004, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9231792999999999, + "after_init_synchronized_e2e_ms_per_call": 0.9443400000000001, + "including_init_host_enqueue_ms_per_call": 4.6166231, + "including_init_synchronized_e2e_ms_per_call": 46.455471900000006, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22318603000000004, + "after_init_synchronized_e2e_ms_per_call": 0.2438196, + "including_init_host_enqueue_ms_per_call": 0.5925304100000001, + "including_init_synchronized_e2e_ms_per_call": 4.794932790000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15318670300000004, + "after_init_synchronized_e2e_ms_per_call": 0.17376756000000002, + "including_init_host_enqueue_ms_per_call": 0.19012114100000002, + "including_init_synchronized_e2e_ms_per_call": 0.6288788790000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.923112, + "synchronized_e2e_ms": 7.949544, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 8.362601, + "median": 0.053824, + "min": 0.04512, + "p90": 0.082464, + "sample_count": 7071 + }, + "host_enqueue_ms": { + "max": 39.889322, + "median": 0.145409, + "min": 0.123456, + "p90": 0.327712, + "sample_count": 7071 + }, + "sample_count": 7071, + "synchronized_e2e_ms": { + "max": 39.926922, + "median": 0.165984, + "min": 0.139425, + "p90": 0.373056, + "sample_count": 7071 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 8401, + "candidate_precomputed_gpu_span_ms": 0.013088, + "candidate_precomputed_host_enqueue_ms": 0.041216, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.013088, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.041216, + "candidate_precomputed_synchronized_e2e_ms": 0.053344, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.013088 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.027392, + "synchronized_e2e_ms": 0.0392 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.013088 + }, + "host_enqueue_ms": { + "median": 0.041216 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.013088 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 8401, + "submission_ms": { + "median": 0.041216 + }, + "synchronized_e2e_ms": { + "median": 0.053344 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7d24c260", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7d24dd90" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.1320293398533006, + "submission": 1.076086956521739, + "synchronized_e2e": 1.1247750449910017 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 7071, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.014816, + "candidate_public_raw_host_enqueue_ms": 0.044352, + "candidate_public_raw_inter_kernel_gap_ms": 0.000192, + "candidate_public_raw_kernel_sum_ms": 0.014592, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.044352, + "candidate_public_raw_synchronized_e2e_ms": 0.06, + "candidate_public_raw_tflops_from_gpu_span": 2.5478358531317498, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.014592 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.204, + "synchronized_e2e_ms": 0.225472 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.014816 + }, + "host_enqueue_ms": { + "median": 0.044352 + }, + "inter_kernel_gap_ms": { + "median": 0.000192 + }, + "kernel_sum_ms": { + "median": 0.014592 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 7071, + "submission_ms": { + "median": 0.044352 + }, + "synchronized_e2e_ms": { + "median": 0.06 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.096002, + "after_init_synchronized_e2e_ms_per_call": 2.120002, + "including_init_host_enqueue_ms_per_call": 39.454024999999994, + "including_init_synchronized_e2e_ms_per_call": 39.555625, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.249517, + "after_init_synchronized_e2e_ms_per_call": 0.2660002, + "including_init_host_enqueue_ms_per_call": 3.9853192999999996, + "including_init_synchronized_e2e_ms_per_call": 4.0095624999999995, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06486850000000001, + "after_init_synchronized_e2e_ms_per_call": 0.08060002, + "including_init_host_enqueue_ms_per_call": 0.43844872999999995, + "including_init_synchronized_e2e_ms_per_call": 0.45495624999999995, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.04640365, + "after_init_synchronized_e2e_ms_per_call": 0.062060001999999996, + "including_init_host_enqueue_ms_per_call": 0.08376167300000001, + "including_init_synchronized_e2e_ms_per_call": 0.09949562499999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.096002, + "synchronized_e2e_ms": 2.120002, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.015392, + "median": 0.014816, + "min": 0.013984, + "p90": 0.015008, + "sample_count": 7071 + }, + "host_enqueue_ms": { + "max": 25.620955, + "median": 0.044352, + "min": 0.035264, + "p90": 0.11888, + "sample_count": 7071 + }, + "sample_count": 7071, + "synchronized_e2e_ms": { + "max": 34.708516, + "median": 0.06, + "min": 0.05152, + "p90": 0.159008, + "sample_count": 7071 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.057185, + "submission_ms": 0.057185, + "synchronized_e2e_ms": 0.072321 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.038304, + "submission_ms": 0.038304, + "synchronized_e2e_ms": 0.052672 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.198177, + "submission_ms": 0.198177, + "synchronized_e2e_ms": 0.216673 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.923112, + "submission_ms": 7.923112, + "synchronized_e2e_ms": 7.949544 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.050144, + "submission_ms": 0.050144, + "synchronized_e2e_ms": 0.065984 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.027392, + "submission_ms": 0.027392, + "synchronized_e2e_ms": 0.0392 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.127745, + "submission_ms": 1.127745, + "synchronized_e2e_ms": 1.147041 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.095105, + "submission_ms": 1.095105, + "synchronized_e2e_ms": 1.114401 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.204, + "submission_ms": 0.204, + "synchronized_e2e_ms": 0.225472 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.096002, + "submission_ms": 2.096002, + "synchronized_e2e_ms": 2.120002 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.581567, + "evolution_kernel_ms": 0.171104, + "evolution_speedup": 3.3989, + "evolution_tflops": 0.2206, + "expected_route": "d288_parent_splitk_hybrid_20260629_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d288_b1_n256_k256_d288", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 30944, + "measurement_schedule_sha256": "d6777fbad0690077aa13b8e4314cf63f04d1222c23f8627806732997b31d0675", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 8401, + "public_pair_count": 7071, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 8401, + "baseline_public_raw": 7071, + "candidate_precomputed": 8401, + "candidate_public_raw": 7071 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 6192 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:post_d895_d288_b1_n256_k256_d288", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.9119804400977995, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.7664, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.749781368130785, + "including_init_synchronized_e2e_speedup": 11.706574298851303, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 3.5501477066558595, + "including_init_synchronized_e2e_speedup": 11.586169787850922, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 3.025056321325975, + "including_init_synchronized_e2e_speedup": 10.539327221024003, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.799992819851988, + "including_init_synchronized_e2e_speedup": 6.320668662566822, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 3.632829373650108, + "hot_synchronized_e2e_speedup": 2.7664, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 228804, + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d288_b1_n256_k256_d288", + "source": "medium_small_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 288, + "K": 8192, + "N": 512, + "baseline_07cf_adapter_bench_iters": 2589, + "baseline_07cf_adapter_gpu_span_ms": 0.242688, + "baseline_07cf_adapter_host_enqueue_ms": 0.14688, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.029888, + "baseline_07cf_adapter_kernel_sum_ms": 0.2128, + "baseline_07cf_adapter_submission_ms": 0.14688, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.343808, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.2128 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.205088, + "synchronized_e2e_ms": 0.390785 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.242688 + }, + "host_enqueue_ms": { + "median": 0.14688 + }, + "inter_kernel_gap_ms": { + "median": 0.029888 + }, + "kernel_sum_ms": { + "median": 0.2128 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2589, + "submission_ms": { + "median": 0.14688 + }, + "synchronized_e2e_ms": { + "median": 0.343808 + } + }, + "baseline_07cf_precomputed_bench_iters": 3805, + "baseline_07cf_precomputed_gpu_span_ms": 0.2376, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042816, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.2376, + "baseline_07cf_precomputed_submission_ms": 0.042816, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.287585, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.2376 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.043264, + "synchronized_e2e_ms": 0.230016 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.2376 + }, + "host_enqueue_ms": { + "median": 0.042816 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.2376 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3805, + "submission_ms": { + "median": 0.042816 + }, + "synchronized_e2e_ms": { + "median": 0.287585 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.0214141414141413, + "submission": 3.4304932735426013, + "synchronized_e2e": 1.1955004607333486 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.607656, + "after_init_synchronized_e2e_ms_per_call": 7.78292, + "including_init_host_enqueue_ms_per_call": 41.835563, + "including_init_synchronized_e2e_ms_per_call": 42.128363, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8929576000000001, + "after_init_synchronized_e2e_ms_per_call": 1.0877192, + "including_init_host_enqueue_ms_per_call": 4.3157483, + "including_init_synchronized_e2e_ms_per_call": 4.522263499999999, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22148776, + "after_init_synchronized_e2e_ms_per_call": 0.4181991199999999, + "including_init_host_enqueue_ms_per_call": 0.56376683, + "including_init_synchronized_e2e_ms_per_call": 0.76165355, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.154340776, + "after_init_synchronized_e2e_ms_per_call": 0.35124711200000003, + "including_init_host_enqueue_ms_per_call": 0.18856868300000001, + "including_init_synchronized_e2e_ms_per_call": 0.385592555, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.607656, + "synchronized_e2e_ms": 7.78292, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.431046, + "median": 0.242688, + "min": 0.23456, + "p90": 0.258816, + "sample_count": 2589 + }, + "host_enqueue_ms": { + "max": 70.265737, + "median": 0.14688, + "min": 0.127584, + "p90": 0.1847744, + "sample_count": 2589 + }, + "sample_count": 2589, + "synchronized_e2e_ms": { + "max": 75.439598, + "median": 0.343808, + "min": 0.326113, + "p90": 0.3784392000000001, + "sample_count": 2589 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3805, + "candidate_precomputed_gpu_span_ms": 0.026624, + "candidate_precomputed_host_enqueue_ms": 0.050144, + "candidate_precomputed_inter_kernel_gap_ms": 0.002272, + "candidate_precomputed_kernel_sum_ms": 0.024128, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.050144, + "candidate_precomputed_synchronized_e2e_ms": 0.063168, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.024128 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.041216, + "synchronized_e2e_ms": 0.057504 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.026624 + }, + "host_enqueue_ms": { + "median": 0.050144 + }, + "inter_kernel_gap_ms": { + "median": 0.002272 + }, + "kernel_sum_ms": { + "median": 0.024128 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3805, + "submission_ms": { + "median": 0.050144 + }, + "synchronized_e2e_ms": { + "median": 0.063168 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb9e51fe60", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb9e51ed80" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.4567307692307694, + "submission": 0.9668155711550734, + "synchronized_e2e": 1.4037487335359675 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2589, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.038784, + "candidate_public_raw_host_enqueue_ms": 0.04848, + "candidate_public_raw_inter_kernel_gap_ms": 9.6e-05, + "candidate_public_raw_kernel_sum_ms": 0.038688, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.04848, + "candidate_public_raw_synchronized_e2e_ms": 0.088672, + "candidate_public_raw_tflops_from_gpu_span": 62.29164356435644, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.038688 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.2584, + "synchronized_e2e_ms": 0.2848 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.038784 + }, + "host_enqueue_ms": { + "median": 0.04848 + }, + "inter_kernel_gap_ms": { + "median": 9.6e-05 + }, + "kernel_sum_ms": { + "median": 0.038688 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2589, + "submission_ms": { + "median": 0.04848 + }, + "synchronized_e2e_ms": { + "median": 0.088672 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 66.394277, + "after_init_synchronized_e2e_ms_per_call": 66.423621, + "including_init_host_enqueue_ms_per_call": 100.971817, + "including_init_synchronized_e2e_ms_per_call": 514.511413, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 6.683059699999999, + "after_init_synchronized_e2e_ms_per_call": 6.7221668999999995, + "including_init_host_enqueue_ms_per_call": 10.140813699999999, + "including_init_synchronized_e2e_ms_per_call": 51.530946099999994, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.71193797, + "after_init_synchronized_e2e_ms_per_call": 0.75202149, + "including_init_host_enqueue_ms_per_call": 1.05771337, + "including_init_synchronized_e2e_ms_per_call": 5.23289941, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.114825797, + "after_init_synchronized_e2e_ms_per_call": 0.155006949, + "including_init_host_enqueue_ms_per_call": 0.149403337, + "including_init_synchronized_e2e_ms_per_call": 0.6030947409999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 66.394277, + "synchronized_e2e_ms": 66.423621, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.03984, + "median": 0.038784, + "min": 0.037696, + "p90": 0.039232, + "sample_count": 2589 + }, + "host_enqueue_ms": { + "max": 5.798374, + "median": 0.04848, + "min": 0.038432, + "p90": 0.06479360000000003, + "sample_count": 2589 + }, + "sample_count": 2589, + "synchronized_e2e_ms": { + "max": 39.139401, + "median": 0.088672, + "min": 0.08048, + "p90": 0.10343680000000002, + "sample_count": 2589 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.065857, + "submission_ms": 0.065857, + "synchronized_e2e_ms": 0.250977 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.043264, + "submission_ms": 0.043264, + "synchronized_e2e_ms": 0.230016 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.205088, + "submission_ms": 0.205088, + "synchronized_e2e_ms": 0.390785 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.607656, + "submission_ms": 7.607656, + "synchronized_e2e_ms": 7.78292 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.087744, + "submission_ms": 0.087744, + "synchronized_e2e_ms": 0.105952 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.041216, + "submission_ms": 0.041216, + "synchronized_e2e_ms": 0.057504 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.315585, + "submission_ms": 1.315585, + "synchronized_e2e_ms": 1.336161 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.234113, + "submission_ms": 1.234113, + "synchronized_e2e_ms": 1.255393 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.2584, + "submission_ms": 0.2584, + "synchronized_e2e_ms": 0.2848 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 66.394277, + "submission_ms": 66.394277, + "synchronized_e2e_ms": 66.423621 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 57.444236, + "evolution_kernel_ms": 0.273615, + "evolution_speedup": 209.9455, + "evolution_tflops": 8.8296, + "expected_route": "d288_parent_splitk_hybrid_20260629_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d288_b1_n512_k8192_d288", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 12788, + "measurement_schedule_sha256": "e262deea11260fd9a983310111212cb371703c880aa91891eec7c3705c7e64f1", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3805, + "public_pair_count": 2589, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3805, + "baseline_public_raw": 2589, + "candidate_precomputed": 3805, + "candidate_public_raw": 2589 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2558 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:post_d895_d288_b1_n512_k8192_d288", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 8.924278846153847, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 3.8773006134969323, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.11717096844208479, + "including_init_synchronized_e2e_speedup": 0.08188032750208402, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.16181079943135598, + "including_init_synchronized_e2e_speedup": 0.08775820826623636, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.5560999593243007, + "including_init_synchronized_e2e_speedup": 0.1455509633042994, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.266008809708267, + "including_init_synchronized_e2e_speedup": 0.6393565202718291, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 6.257425742574258, + "hot_synchronized_e2e_speedup": 3.8773006134969323, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 228803, + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d288_b1_n512_k8192_d288", + "source": "high_k_low_n", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 288, + "K": 1024, + "N": 2048, + "baseline_07cf_adapter_bench_iters": 2340, + "baseline_07cf_adapter_gpu_span_ms": 0.073344, + "baseline_07cf_adapter_host_enqueue_ms": 0.14726450000000002, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.034544000000000005, + "baseline_07cf_adapter_kernel_sum_ms": 0.038784, + "baseline_07cf_adapter_submission_ms": 0.14726450000000002, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.174592, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.038784 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.301344, + "synchronized_e2e_ms": 0.324448 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.073344 + }, + "host_enqueue_ms": { + "median": 0.14726450000000002 + }, + "inter_kernel_gap_ms": { + "median": 0.034544000000000005 + }, + "kernel_sum_ms": { + "median": 0.038784 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2340, + "submission_ms": { + "median": 0.14726450000000002 + }, + "synchronized_e2e_ms": { + "median": 0.174592 + } + }, + "baseline_07cf_precomputed_bench_iters": 3030, + "baseline_07cf_precomputed_gpu_span_ms": 0.032992, + "baseline_07cf_precomputed_host_enqueue_ms": 0.0431685, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.032992, + "baseline_07cf_precomputed_submission_ms": 0.0431685, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.08352, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.032992 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.046208, + "synchronized_e2e_ms": 0.078528 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.032992 + }, + "host_enqueue_ms": { + "median": 0.0431685 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.032992 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3030, + "submission_ms": { + "median": 0.0431685 + }, + "synchronized_e2e_ms": { + "median": 0.08352 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.223084384093114, + "submission": 3.4113879333310173, + "synchronized_e2e": 2.0904214559386975 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.546825, + "after_init_synchronized_e2e_ms_per_call": 8.574601, + "including_init_host_enqueue_ms_per_call": 42.996269, + "including_init_synchronized_e2e_ms_per_call": 492.616095, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.98722055, + "after_init_synchronized_e2e_ms_per_call": 1.0145928999999998, + "including_init_host_enqueue_ms_per_call": 4.43216495, + "including_init_synchronized_e2e_ms_per_call": 49.4187423, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.231260105, + "after_init_synchronized_e2e_ms_per_call": 0.25859209, + "including_init_host_enqueue_ms_per_call": 0.575754545, + "including_init_synchronized_e2e_ms_per_call": 5.099007029999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15566406050000003, + "after_init_synchronized_e2e_ms_per_call": 0.18299200899999998, + "including_init_host_enqueue_ms_per_call": 0.19011350450000003, + "including_init_synchronized_e2e_ms_per_call": 0.667033503, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.546825, + "synchronized_e2e_ms": 8.574601, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.548641, + "median": 0.073344, + "min": 0.064128, + "p90": 0.0927721, + "sample_count": 2340 + }, + "host_enqueue_ms": { + "max": 0.978721, + "median": 0.14726450000000002, + "min": 0.122944, + "p90": 0.18806489999999998, + "sample_count": 2340 + }, + "sample_count": 2340, + "synchronized_e2e_ms": { + "max": 1.131105, + "median": 0.174592, + "min": 0.151296, + "p90": 0.2147712, + "sample_count": 2340 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3030, + "candidate_precomputed_gpu_span_ms": 0.037952, + "candidate_precomputed_host_enqueue_ms": 0.040448, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.037952, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.040448, + "candidate_precomputed_synchronized_e2e_ms": 0.0736, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.037952 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.035552, + "synchronized_e2e_ms": 0.059232 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.037952 + }, + "host_enqueue_ms": { + "median": 0.040448 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.037952 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3030, + "submission_ms": { + "median": 0.040448 + }, + "synchronized_e2e_ms": { + "median": 0.0736 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb722c6ba0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb722c72f0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.1281618887015177, + "submission": 1.1115629944620253, + "synchronized_e2e": 1.2243478260869565 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 2340, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.042816, + "candidate_public_raw_host_enqueue_ms": 0.0449605, + "candidate_public_raw_inter_kernel_gap_ms": 0.00016, + "candidate_public_raw_kernel_sum_ms": 0.042656, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.0449605, + "candidate_public_raw_synchronized_e2e_ms": 0.090112, + "candidate_public_raw_tflops_from_gpu_span": 56.425614349775785, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.042656 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.307841, + "synchronized_e2e_ms": 0.336897 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.042816 + }, + "host_enqueue_ms": { + "median": 0.0449605 + }, + "inter_kernel_gap_ms": { + "median": 0.00016 + }, + "kernel_sum_ms": { + "median": 0.042656 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2340, + "submission_ms": { + "median": 0.0449605 + }, + "synchronized_e2e_ms": { + "median": 0.090112 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 57.730268, + "after_init_synchronized_e2e_ms_per_call": 57.76726, + "including_init_host_enqueue_ms_per_call": 92.604352, + "including_init_synchronized_e2e_ms_per_call": 92.710336, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 5.81349125, + "after_init_synchronized_e2e_ms_per_call": 5.8578268, + "including_init_host_enqueue_ms_per_call": 9.300899650000002, + "including_init_synchronized_e2e_ms_per_call": 9.3521344, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.621813575, + "after_init_synchronized_e2e_ms_per_call": 0.6668834800000001, + "including_init_host_enqueue_ms_per_call": 0.9705544150000001, + "including_init_synchronized_e2e_ms_per_call": 1.01631424, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.1026458075, + "after_init_synchronized_e2e_ms_per_call": 0.147789148, + "including_init_host_enqueue_ms_per_call": 0.1375198915, + "including_init_synchronized_e2e_ms_per_call": 0.182732224, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 57.730268, + "synchronized_e2e_ms": 57.76726, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.04352, + "median": 0.042816, + "min": 0.042464, + "p90": 0.043072, + "sample_count": 2340 + }, + "host_enqueue_ms": { + "max": 10.543722, + "median": 0.0449605, + "min": 0.036288, + "p90": 0.0630496, + "sample_count": 2340 + }, + "sample_count": 2340, + "synchronized_e2e_ms": { + "max": 20.825365, + "median": 0.090112, + "min": 0.082464, + "p90": 0.10604799999999998, + "sample_count": 2340 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.069568, + "submission_ms": 0.069568, + "synchronized_e2e_ms": 0.100064 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.046208, + "submission_ms": 0.046208, + "synchronized_e2e_ms": 0.078528 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.301344, + "submission_ms": 0.301344, + "synchronized_e2e_ms": 0.324448 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.546825, + "submission_ms": 8.546825, + "synchronized_e2e_ms": 8.574601 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.069024, + "submission_ms": 0.069024, + "synchronized_e2e_ms": 0.085792 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.035552, + "submission_ms": 0.035552, + "synchronized_e2e_ms": 0.059232 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.540449, + "submission_ms": 1.540449, + "synchronized_e2e_ms": 1.565121 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.208706, + "submission_ms": 1.208706, + "synchronized_e2e_ms": 1.23005 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.307841, + "submission_ms": 0.307841, + "synchronized_e2e_ms": 0.336897 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 57.730268, + "submission_ms": 57.730268, + "synchronized_e2e_ms": 57.76726 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.65016, + "evolution_kernel_ms": 0.18616, + "evolution_speedup": 3.4925, + "evolution_tflops": 12.9777, + "expected_route": "d288_parent_splitk_hybrid_20260629_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d288_b2_n2048_k1024_d288", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 10740, + "measurement_schedule_sha256": "9d25518976dd1558e266b450781f5e00dc2229e6482edd23c3dd9b9d00ae1e4e", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3030, + "public_pair_count": 2340, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3030, + "baseline_public_raw": 2340, + "candidate_precomputed": 3030, + "candidate_public_raw": 2340 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2148 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:post_d895_d288_b2_n2048_k1024_d288", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.8693086003372682, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.9375, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.14843357638911728, + "including_init_synchronized_e2e_speedup": 5.313497030147749, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.17320295301322325, + "including_init_synchronized_e2e_speedup": 5.2842207122258635, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.3877620270335681, + "including_init_synchronized_e2e_speedup": 5.017155943815172, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.2381965217094286, + "including_init_synchronized_e2e_speedup": 3.6503331946531774, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.7130044843049328, + "hot_synchronized_e2e_speedup": 1.9375, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 228801, + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d288_b2_n2048_k1024_d288", + "source": "new_medium_gap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 288, + "K": 1024, + "N": 8192, + "baseline_07cf_adapter_bench_iters": 889, + "baseline_07cf_adapter_gpu_span_ms": 0.167808, + "baseline_07cf_adapter_host_enqueue_ms": 0.169217, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.004128, + "baseline_07cf_adapter_kernel_sum_ms": 0.163489, + "baseline_07cf_adapter_submission_ms": 0.169217, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.280832, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.163489 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.218336, + "synchronized_e2e_ms": 0.32368 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.167808 + }, + "host_enqueue_ms": { + "median": 0.169217 + }, + "inter_kernel_gap_ms": { + "median": 0.004128 + }, + "kernel_sum_ms": { + "median": 0.163489 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 889, + "submission_ms": { + "median": 0.169217 + }, + "synchronized_e2e_ms": { + "median": 0.280832 + } + }, + "baseline_07cf_precomputed_bench_iters": 1320, + "baseline_07cf_precomputed_gpu_span_ms": 0.1246405, + "baseline_07cf_precomputed_host_enqueue_ms": 0.05152, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.1246405, + "baseline_07cf_precomputed_submission_ms": 0.05152, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.182624, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.1246405 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.039968, + "synchronized_e2e_ms": 0.16384 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.1246405 + }, + "host_enqueue_ms": { + "median": 0.05152 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.1246405 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1320, + "submission_ms": { + "median": 0.05152 + }, + "synchronized_e2e_ms": { + "median": 0.182624 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.3463360625157954, + "submission": 3.284491459627329, + "synchronized_e2e": 1.5377606448221484 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.820873, + "after_init_synchronized_e2e_ms_per_call": 8.898218, + "including_init_host_enqueue_ms_per_call": 44.372141, + "including_init_synchronized_e2e_ms_per_call": 44.529455, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 1.0343826, + "after_init_synchronized_e2e_ms_per_call": 1.1425706, + "including_init_host_enqueue_ms_per_call": 4.5895094, + "including_init_synchronized_e2e_ms_per_call": 4.705694299999999, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.25573356, + "after_init_synchronized_e2e_ms_per_call": 0.36700586, + "including_init_host_enqueue_ms_per_call": 0.6112462399999999, + "including_init_synchronized_e2e_ms_per_call": 0.72331823, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.177868656, + "after_init_synchronized_e2e_ms_per_call": 0.28944938600000003, + "including_init_host_enqueue_ms_per_call": 0.213419924, + "including_init_synchronized_e2e_ms_per_call": 0.325080623, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.820873, + "synchronized_e2e_ms": 8.898218, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.420512, + "median": 0.167808, + "min": 0.157217, + "p90": 0.176768, + "sample_count": 889 + }, + "host_enqueue_ms": { + "max": 0.437248, + "median": 0.169217, + "min": 0.145728, + "p90": 0.19668480000000002, + "sample_count": 889 + }, + "sample_count": 889, + "synchronized_e2e_ms": { + "max": 0.535169, + "median": 0.280832, + "min": 0.261376, + "p90": 0.3058368, + "sample_count": 889 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 1320, + "candidate_precomputed_gpu_span_ms": 0.075616, + "candidate_precomputed_host_enqueue_ms": 0.046368, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.075616, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.046368, + "candidate_precomputed_synchronized_e2e_ms": 0.11611250000000001, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.075616 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.028064, + "synchronized_e2e_ms": 0.079712 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.075616 + }, + "host_enqueue_ms": { + "median": 0.046368 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.075616 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1320, + "submission_ms": { + "median": 0.046368 + }, + "synchronized_e2e_ms": { + "median": 0.11611250000000001 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb933d6e70", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb933d64b0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.487092678798138, + "submission": 1.2318840579710144, + "synchronized_e2e": 1.47085800409086 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 889, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.112448, + "candidate_public_raw_host_enqueue_ms": 0.05712, + "candidate_public_raw_inter_kernel_gap_ms": 0.00016, + "candidate_public_raw_kernel_sum_ms": 0.112288, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.05712, + "candidate_public_raw_synchronized_e2e_ms": 0.170785, + "candidate_public_raw_tflops_from_gpu_span": 171.87813773477518, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.112288 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.284768, + "synchronized_e2e_ms": 0.370433 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.112448 + }, + "host_enqueue_ms": { + "median": 0.05712 + }, + "inter_kernel_gap_ms": { + "median": 0.00016 + }, + "kernel_sum_ms": { + "median": 0.112288 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 889, + "submission_ms": { + "median": 0.05712 + }, + "synchronized_e2e_ms": { + "median": 0.170785 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 43.710573, + "after_init_synchronized_e2e_ms_per_call": 43.776109, + "including_init_host_enqueue_ms_per_call": 79.548593, + "including_init_synchronized_e2e_ms_per_call": 527.876322, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 4.4224653, + "after_init_synchronized_e2e_ms_per_call": 4.5313174, + "including_init_host_enqueue_ms_per_call": 8.006267300000001, + "including_init_synchronized_e2e_ms_per_call": 52.941338699999996, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.4936545299999999, + "after_init_synchronized_e2e_ms_per_call": 0.60683824, + "including_init_host_enqueue_ms_per_call": 0.8520347300000001, + "including_init_synchronized_e2e_ms_per_call": 5.44784037, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.10077345299999999, + "after_init_synchronized_e2e_ms_per_call": 0.214390324, + "including_init_host_enqueue_ms_per_call": 0.13661147299999998, + "including_init_synchronized_e2e_ms_per_call": 0.6984905369999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 43.710573, + "synchronized_e2e_ms": 43.776109, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.11328, + "median": 0.112448, + "min": 0.111648, + "p90": 0.11280720000000001, + "sample_count": 889 + }, + "host_enqueue_ms": { + "max": 0.120224, + "median": 0.05712, + "min": 0.045152, + "p90": 0.0665088, + "sample_count": 889 + }, + "sample_count": 889, + "synchronized_e2e_ms": { + "max": 0.226496, + "median": 0.170785, + "min": 0.159584, + "p90": 0.17917519999999998, + "sample_count": 889 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.070528, + "submission_ms": 0.070528, + "synchronized_e2e_ms": 0.176064 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.039968, + "submission_ms": 0.039968, + "synchronized_e2e_ms": 0.16384 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.218336, + "submission_ms": 0.218336, + "synchronized_e2e_ms": 0.32368 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.820873, + "submission_ms": 8.820873, + "synchronized_e2e_ms": 8.898218 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.074688, + "submission_ms": 0.074688, + "synchronized_e2e_ms": 0.112352 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.028064, + "submission_ms": 0.028064, + "synchronized_e2e_ms": 0.079712 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.186305, + "submission_ms": 1.186305, + "synchronized_e2e_ms": 1.208161 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.011041, + "submission_ms": 1.011041, + "synchronized_e2e_ms": 1.031233 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.284768, + "submission_ms": 0.284768, + "synchronized_e2e_ms": 0.370433 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 43.710573, + "submission_ms": 43.710573, + "synchronized_e2e_ms": 43.776109 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.669808, + "evolution_kernel_ms": 0.214687, + "evolution_speedup": 3.1199, + "evolution_tflops": 90.0257, + "expected_route": "d288_parent_splitk_hybrid_20260629_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d288_b4_n8192_k1024_d288", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 4418, + "measurement_schedule_sha256": "00db8c0178d1d06f9eb57922ddcd4ea61d1c7c4b8277b5943b7749c928cd2749", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 1320, + "public_pair_count": 889, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 1320, + "baseline_public_raw": 889, + "candidate_precomputed": 1320, + "candidate_public_raw": 889 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 884 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:post_d895_d288_b4_n8192_k1024_d288", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.648335008463817, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.6443598676698776, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.20326653517789808, + "including_init_synchronized_e2e_speedup": 0.08435584841405333, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.2521497611268635, + "including_init_synchronized_e2e_speedup": 0.08888506440431208, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.6047836734876827, + "including_init_synchronized_e2e_speedup": 0.13277155365695856, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.3501047090166254, + "including_init_synchronized_e2e_speedup": 0.4654044769113316, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.4923164484917473, + "hot_synchronized_e2e_speedup": 1.6443598676698776, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 228802, + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d288_b4_n8192_k1024_d288", + "source": "new_medium_gap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 32, + "K": 1024, + "N": 32768, + "baseline_07cf_adapter_bench_iters": 783, + "baseline_07cf_adapter_gpu_span_ms": 0.143552, + "baseline_07cf_adapter_host_enqueue_ms": 0.146432, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.002176, + "baseline_07cf_adapter_kernel_sum_ms": 0.141376, + "baseline_07cf_adapter_submission_ms": 0.146432, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.24096, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.141376 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.204289, + "synchronized_e2e_ms": 0.266017 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.143552 + }, + "host_enqueue_ms": { + "median": 0.146432 + }, + "inter_kernel_gap_ms": { + "median": 0.002176 + }, + "kernel_sum_ms": { + "median": 0.141376 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 783, + "submission_ms": { + "median": 0.146432 + }, + "synchronized_e2e_ms": { + "median": 0.24096 + } + }, + "baseline_07cf_precomputed_bench_iters": 1611, + "baseline_07cf_precomputed_gpu_span_ms": 0.070912, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042112, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.070912, + "baseline_07cf_precomputed_submission_ms": 0.042112, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.118816, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.070912 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.04192, + "synchronized_e2e_ms": 0.115904 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.070912 + }, + "host_enqueue_ms": { + "median": 0.042112 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.070912 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1611, + "submission_ms": { + "median": 0.042112 + }, + "synchronized_e2e_ms": { + "median": 0.118816 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.0243682310469313, + "submission": 3.4772036474164136, + "synchronized_e2e": 2.028009695663884 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.912328, + "after_init_synchronized_e2e_ms_per_call": 7.962696, + "including_init_host_enqueue_ms_per_call": 44.846766, + "including_init_synchronized_e2e_ms_per_call": 463.07401500000003, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9230215999999999, + "after_init_synchronized_e2e_ms_per_call": 1.0131336000000002, + "including_init_host_enqueue_ms_per_call": 4.6164654, + "including_init_synchronized_e2e_ms_per_call": 46.5242655, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22409096, + "after_init_synchronized_e2e_ms_per_call": 0.31817736, + "including_init_host_enqueue_ms_per_call": 0.5934353400000001, + "including_init_synchronized_e2e_ms_per_call": 4.8692905500000006, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.154197896, + "after_init_synchronized_e2e_ms_per_call": 0.24868173600000001, + "including_init_host_enqueue_ms_per_call": 0.19113233400000001, + "including_init_synchronized_e2e_ms_per_call": 0.7037930550000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.912328, + "synchronized_e2e_ms": 7.962696, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.41456, + "median": 0.143552, + "min": 0.141856, + "p90": 0.14448, + "sample_count": 783 + }, + "host_enqueue_ms": { + "max": 0.456704, + "median": 0.146432, + "min": 0.129472, + "p90": 0.16629760000000002, + "sample_count": 783 + }, + "sample_count": 783, + "synchronized_e2e_ms": { + "max": 0.509024, + "median": 0.24096, + "min": 0.231937, + "p90": 0.2523842, + "sample_count": 783 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 1611, + "candidate_precomputed_gpu_span_ms": 0.058944, + "candidate_precomputed_host_enqueue_ms": 0.052992, + "candidate_precomputed_inter_kernel_gap_ms": 0.001856, + "candidate_precomputed_kernel_sum_ms": 0.056992, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.052992, + "candidate_precomputed_synchronized_e2e_ms": 0.09232, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.056992 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.040064, + "synchronized_e2e_ms": 0.08416 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.058944 + }, + "host_enqueue_ms": { + "median": 0.052992 + }, + "inter_kernel_gap_ms": { + "median": 0.001856 + }, + "kernel_sum_ms": { + "median": 0.056992 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1611, + "submission_ms": { + "median": 0.052992 + }, + "synchronized_e2e_ms": { + "median": 0.09232 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282afc50", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282c6e10" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 2.1672095548317043, + "submission": 0.838768115942029, + "synchronized_e2e": 1.8797227036395148 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 783, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.127744, + "candidate_public_raw_host_enqueue_ms": 0.044448, + "candidate_public_raw_inter_kernel_gap_ms": 0.000128, + "candidate_public_raw_kernel_sum_ms": 0.127617, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.044448, + "candidate_public_raw_synchronized_e2e_ms": 0.173536, + "candidate_public_raw_tflops_from_gpu_span": 67.2433507014028, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.127617 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.051136, + "synchronized_e2e_ms": 0.172928 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.127744 + }, + "host_enqueue_ms": { + "median": 0.044448 + }, + "inter_kernel_gap_ms": { + "median": 0.000128 + }, + "kernel_sum_ms": { + "median": 0.127617 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 783, + "submission_ms": { + "median": 0.044448 + }, + "synchronized_e2e_ms": { + "median": 0.173536 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 69.48084, + "after_init_synchronized_e2e_ms_per_call": 69.596136, + "including_init_host_enqueue_ms_per_call": 106.838863, + "including_init_synchronized_e2e_ms_per_call": 107.031759, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 6.9880872, + "after_init_synchronized_e2e_ms_per_call": 7.1157960000000005, + "including_init_host_enqueue_ms_per_call": 10.7238895, + "including_init_synchronized_e2e_ms_per_call": 10.8593583, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.73881192, + "after_init_synchronized_e2e_ms_per_call": 0.867762, + "including_init_host_enqueue_ms_per_call": 1.11239215, + "including_init_synchronized_e2e_ms_per_call": 1.24211823, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.113884392, + "after_init_synchronized_e2e_ms_per_call": 0.2429586, + "including_init_host_enqueue_ms_per_call": 0.151242415, + "including_init_synchronized_e2e_ms_per_call": 0.280394223, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 69.48084, + "synchronized_e2e_ms": 69.596136, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.128448, + "median": 0.127744, + "min": 0.127392, + "p90": 0.127968, + "sample_count": 783 + }, + "host_enqueue_ms": { + "max": 0.081472, + "median": 0.044448, + "min": 0.037632, + "p90": 0.05214720000000002, + "sample_count": 783 + }, + "sample_count": 783, + "synchronized_e2e_ms": { + "max": 0.20688, + "median": 0.173536, + "min": 0.167488, + "p90": 0.1799616, + "sample_count": 783 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.067392, + "submission_ms": 0.067392, + "synchronized_e2e_ms": 0.140128 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.04192, + "submission_ms": 0.04192, + "synchronized_e2e_ms": 0.115904 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.204289, + "submission_ms": 0.204289, + "synchronized_e2e_ms": 0.266017 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.912328, + "submission_ms": 7.912328, + "synchronized_e2e_ms": 7.962696 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.070112, + "submission_ms": 0.070112, + "synchronized_e2e_ms": 0.102976 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.040064, + "submission_ms": 0.040064, + "synchronized_e2e_ms": 0.08416 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.151969, + "submission_ms": 1.151969, + "synchronized_e2e_ms": 1.178529 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.586178, + "submission_ms": 1.586178, + "synchronized_e2e_ms": 1.609058 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.051136, + "submission_ms": 0.051136, + "synchronized_e2e_ms": 0.172928 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 69.48084, + "submission_ms": 69.48084, + "synchronized_e2e_ms": 69.596136 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.360896, + "evolution_kernel_ms": 0.218816, + "evolution_speedup": 1.6493, + "evolution_tflops": 39.2564, + "expected_route": "microdim_hybrid_9c0d_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d32_b4_n32768_k1024_d32", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 4788, + "measurement_schedule_sha256": "8ca5874a0cdb6f0efa2ef15399b9cf44e8be28fd843ca8779328212da76a3a15", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 1611, + "public_pair_count": 783, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 1611, + "baseline_public_raw": 783, + "candidate_precomputed": 1611, + "candidate_public_raw": 783 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 960 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:post_d895_d32_b4_n32768_k1024_d32", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.203040173724213, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.3885303337635995, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.11441290361292472, + "including_init_synchronized_e2e_speedup": 4.3265103678245636, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.1423781120200748, + "including_init_synchronized_e2e_speedup": 4.28425549785939, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.36666431579165715, + "including_init_synchronized_e2e_speedup": 3.9201506204445615, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.0235560132467014, + "including_init_synchronized_e2e_speedup": 2.510012679540834, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.1237474949899802, + "hot_synchronized_e2e_speedup": 1.3885303337635995, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 23202, + "selected_route": "microdim_hybrid_9c0d_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d32_b4_n32768_k1024_d32", + "source": "near_floor_microdim", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 32, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 8, + "D": 32, + "K": 512, + "N": 65536, + "baseline_07cf_adapter_bench_iters": 265, + "baseline_07cf_adapter_gpu_span_ms": 0.378369, + "baseline_07cf_adapter_host_enqueue_ms": 0.14912, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.00208, + "baseline_07cf_adapter_kernel_sum_ms": 0.376224, + "baseline_07cf_adapter_submission_ms": 0.14912, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.479329, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.376224 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.246944, + "synchronized_e2e_ms": 0.50768 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.378369 + }, + "host_enqueue_ms": { + "median": 0.14912 + }, + "inter_kernel_gap_ms": { + "median": 0.00208 + }, + "kernel_sum_ms": { + "median": 0.376224 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 265, + "submission_ms": { + "median": 0.14912 + }, + "synchronized_e2e_ms": { + "median": 0.479329 + } + }, + "baseline_07cf_precomputed_bench_iters": 984, + "baseline_07cf_precomputed_gpu_span_ms": 0.101664, + "baseline_07cf_precomputed_host_enqueue_ms": 0.04224, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.101664, + "baseline_07cf_precomputed_submission_ms": 0.04224, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.150496, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.101664 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.047008, + "synchronized_e2e_ms": 0.150624 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.101664 + }, + "host_enqueue_ms": { + "median": 0.04224 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.101664 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 984, + "submission_ms": { + "median": 0.04224 + }, + "synchronized_e2e_ms": { + "median": 0.150496 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 3.721759915014164, + "submission": 3.5303030303030303, + "synchronized_e2e": 3.1849949500318946 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.241224, + "after_init_synchronized_e2e_ms_per_call": 7.314888, + "including_init_host_enqueue_ms_per_call": 41.469131000000004, + "including_init_synchronized_e2e_ms_per_call": 41.660331, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8583304, + "after_init_synchronized_e2e_ms_per_call": 1.1628848999999999, + "including_init_host_enqueue_ms_per_call": 4.281121100000001, + "including_init_synchronized_e2e_ms_per_call": 4.5974292, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22004104000000002, + "after_init_synchronized_e2e_ms_per_call": 0.54768459, + "including_init_host_enqueue_ms_per_call": 0.56232011, + "including_init_synchronized_e2e_ms_per_call": 0.89113902, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.156212104, + "after_init_synchronized_e2e_ms_per_call": 0.486164559, + "including_init_host_enqueue_ms_per_call": 0.190440011, + "including_init_synchronized_e2e_ms_per_call": 0.520510002, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.241224, + "synchronized_e2e_ms": 7.314888, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.38016, + "median": 0.378369, + "min": 0.376801, + "p90": 0.3792842, + "sample_count": 265 + }, + "host_enqueue_ms": { + "max": 0.277729, + "median": 0.14912, + "min": 0.131968, + "p90": 0.1729216, + "sample_count": 265 + }, + "sample_count": 265, + "synchronized_e2e_ms": { + "max": 0.555521, + "median": 0.479329, + "min": 0.471137, + "p90": 0.4957572, + "sample_count": 265 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 2, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 984, + "candidate_precomputed_gpu_span_ms": 0.115872, + "candidate_precomputed_host_enqueue_ms": 0.039792, + "candidate_precomputed_inter_kernel_gap_ms": 0.0, + "candidate_precomputed_kernel_sum_ms": 0.115872, + "candidate_precomputed_launch_count": 1, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.039792, + "candidate_precomputed_synchronized_e2e_ms": 0.149952, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.115872 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.040032, + "synchronized_e2e_ms": 0.144032 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.115872 + }, + "host_enqueue_ms": { + "median": 0.039792 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.115872 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 984, + "submission_ms": { + "median": 0.039792 + }, + "synchronized_e2e_ms": { + "median": 0.149952 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb9e4d4260", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb9e4d4920" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 3.3631593482463407, + "submission": 1.174909529553679, + "synchronized_e2e": 2.9223218096457533 + }, + "candidate_public_raw_assignment_launch_count": 1, + "candidate_public_raw_bench_iters": 265, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.389696, + "candidate_public_raw_host_enqueue_ms": 0.046752, + "candidate_public_raw_inter_kernel_gap_ms": 0.000192, + "candidate_public_raw_kernel_sum_ms": 0.389504, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.046752, + "candidate_public_raw_synchronized_e2e_ms": 0.438208, + "candidate_public_raw_tflops_from_gpu_span": 44.08531056002627, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.389504 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.294336, + "synchronized_e2e_ms": 0.672641 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.389696 + }, + "host_enqueue_ms": { + "median": 0.046752 + }, + "inter_kernel_gap_ms": { + "median": 0.000192 + }, + "kernel_sum_ms": { + "median": 0.389504 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 265, + "submission_ms": { + "median": 0.046752 + }, + "synchronized_e2e_ms": { + "median": 0.438208 + } + }, + "candidate_public_raw_total_launch_count": 2, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 42.274252, + "after_init_synchronized_e2e_ms_per_call": 42.638444, + "including_init_host_enqueue_ms_per_call": 76.85179199999999, + "including_init_synchronized_e2e_ms_per_call": 490.726236, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 4.269502, + "after_init_synchronized_e2e_ms_per_call": 4.6582316, + "including_init_host_enqueue_ms_per_call": 7.727255999999999, + "including_init_synchronized_e2e_ms_per_call": 49.4670108, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.46902699999999997, + "after_init_synchronized_e2e_ms_per_call": 0.86021036, + "including_init_host_enqueue_ms_per_call": 0.8148023999999999, + "including_init_synchronized_e2e_ms_per_call": 5.34108828, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.0889795, + "after_init_synchronized_e2e_ms_per_call": 0.480408236, + "including_init_host_enqueue_ms_per_call": 0.12355704, + "including_init_synchronized_e2e_ms_per_call": 0.928496028, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 42.274252, + "synchronized_e2e_ms": 42.638444, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.390465, + "median": 0.389696, + "min": 0.38912, + "p90": 0.39004859999999997, + "sample_count": 265 + }, + "host_enqueue_ms": { + "max": 0.079712, + "median": 0.046752, + "min": 0.03968, + "p90": 0.057273599999999994, + "sample_count": 265 + }, + "sample_count": 265, + "synchronized_e2e_ms": { + "max": 0.465568, + "median": 0.438208, + "min": 0.432257, + "p90": 0.44667559999999995, + "sample_count": 265 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.06896, + "submission_ms": 0.06896, + "synchronized_e2e_ms": 0.170944 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.047008, + "submission_ms": 0.047008, + "synchronized_e2e_ms": 0.150624 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.246944, + "submission_ms": 0.246944, + "synchronized_e2e_ms": 0.50768 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.241224, + "submission_ms": 7.241224, + "synchronized_e2e_ms": 7.314888 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.058784, + "submission_ms": 0.058784, + "synchronized_e2e_ms": 0.160352 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.040032, + "submission_ms": 0.040032, + "synchronized_e2e_ms": 0.144032 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.34797, + "submission_ms": 1.34797, + "synchronized_e2e_ms": 1.37005 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.125761, + "submission_ms": 1.125761, + "synchronized_e2e_ms": 1.159681 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.294336, + "submission_ms": 0.294336, + "synchronized_e2e_ms": 0.672641 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 42.274252, + "submission_ms": 42.274252, + "synchronized_e2e_ms": 42.638444 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.500015, + "evolution_kernel_ms": 0.301439, + "evolution_speedup": 1.6588, + "evolution_tflops": 56.9928, + "expected_route": "microdim_pipeline4_08f9_v4", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d32_b8_n65536_k512_d32", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 2498, + "measurement_schedule_sha256": "0859c0b5e0b27740c90cd374696ce687280cef8278e82b43c85fd2b595401ef9", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 984, + "public_pair_count": 265, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 984, + "baseline_public_raw": 265, + "candidate_precomputed": 984, + "candidate_public_raw": 265 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 500 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:post_d895_d32_b8_n65536_k512_d32", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.8773819386909694, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.093838998831605, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.171556166543038, + "including_init_synchronized_e2e_speedup": 0.08489525919702406, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.24964085083274948, + "including_init_synchronized_e2e_speedup": 0.09293929682931236, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.6366868099565786, + "including_init_synchronized_e2e_speedup": 0.1668459634597165, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.0119821488655745, + "including_init_synchronized_e2e_speedup": 0.5605947535620476, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 0.9709337534898999, + "hot_synchronized_e2e_speedup": 1.093838998831605, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 23201, + "selected_route": "microdim_pipeline4_08f9_v4", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d32_b8_n65536_k512_d32", + "source": "near_floor_microdim", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 128, + "D_PAD": 32, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 352, + "K": 256, + "N": 256, + "baseline_07cf_adapter_bench_iters": 6232, + "baseline_07cf_adapter_gpu_span_ms": 0.055072, + "baseline_07cf_adapter_host_enqueue_ms": 0.148256, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.042048, + "baseline_07cf_adapter_kernel_sum_ms": 0.013056, + "baseline_07cf_adapter_submission_ms": 0.148256, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.16857650000000002, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.013056 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.24736, + "synchronized_e2e_ms": 0.270944 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.055072 + }, + "host_enqueue_ms": { + "median": 0.148256 + }, + "inter_kernel_gap_ms": { + "median": 0.042048 + }, + "kernel_sum_ms": { + "median": 0.013056 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 6232, + "submission_ms": { + "median": 0.148256 + }, + "synchronized_e2e_ms": { + "median": 0.16857650000000002 + } + }, + "baseline_07cf_precomputed_bench_iters": 8160, + "baseline_07cf_precomputed_gpu_span_ms": 0.012416, + "baseline_07cf_precomputed_host_enqueue_ms": 0.04352, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.012416, + "baseline_07cf_precomputed_submission_ms": 0.04352, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.063296, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.012416 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.05392, + "synchronized_e2e_ms": 0.071136 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.012416 + }, + "host_enqueue_ms": { + "median": 0.04352 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.012416 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 8160, + "submission_ms": { + "median": 0.04352 + }, + "synchronized_e2e_ms": { + "median": 0.063296 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 4.435567010309279, + "submission": 3.406617647058823, + "synchronized_e2e": 2.6633041582406474 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.196713, + "after_init_synchronized_e2e_ms_per_call": 8.224233, + "including_init_host_enqueue_ms_per_call": 42.646157, + "including_init_synchronized_e2e_ms_per_call": 492.265727, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9531017, + "after_init_synchronized_e2e_ms_per_call": 0.9741421499999999, + "including_init_host_enqueue_ms_per_call": 4.3980461, + "including_init_synchronized_e2e_ms_per_call": 49.37829155, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22874057, + "after_init_synchronized_e2e_ms_per_call": 0.24913306500000004, + "including_init_host_enqueue_ms_per_call": 0.57323501, + "including_init_synchronized_e2e_ms_per_call": 5.089548005, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15630445699999998, + "after_init_synchronized_e2e_ms_per_call": 0.1766321565, + "including_init_host_enqueue_ms_per_call": 0.19075390099999998, + "including_init_synchronized_e2e_ms_per_call": 0.6606736505, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.196713, + "synchronized_e2e_ms": 8.224233, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.301446, + "median": 0.055072, + "min": 0.04544, + "p90": 0.08352960000000012, + "sample_count": 6232 + }, + "host_enqueue_ms": { + "max": 39.220872, + "median": 0.148256, + "min": 0.123616, + "p90": 0.2500809, + "sample_count": 6232 + }, + "sample_count": 6232, + "synchronized_e2e_ms": { + "max": 39.269032, + "median": 0.16857650000000002, + "min": 0.141568, + "p90": 0.2777889, + "sample_count": 6232 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 8160, + "candidate_precomputed_gpu_span_ms": 0.01968, + "candidate_precomputed_host_enqueue_ms": 0.056448, + "candidate_precomputed_inter_kernel_gap_ms": 0.006688, + "candidate_precomputed_kernel_sum_ms": 0.013024, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.056448, + "candidate_precomputed_synchronized_e2e_ms": 0.069184, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.013024 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.050624, + "synchronized_e2e_ms": 0.067136 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.01968 + }, + "host_enqueue_ms": { + "median": 0.056448 + }, + "inter_kernel_gap_ms": { + "median": 0.006688 + }, + "kernel_sum_ms": { + "median": 0.013024 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 8160, + "submission_ms": { + "median": 0.056448 + }, + "synchronized_e2e_ms": { + "median": 0.069184 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7eda3050", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7eda2480" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 0.8357723577235773, + "submission": 0.8185941043083901, + "synchronized_e2e": 0.9458834412580944 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 6232, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.016448, + "candidate_public_raw_host_enqueue_ms": 0.046208, + "candidate_public_raw_inter_kernel_gap_ms": 3.2e-05, + "candidate_public_raw_kernel_sum_ms": 0.016544, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.046208, + "candidate_public_raw_synchronized_e2e_ms": 0.06544, + "candidate_public_raw_tflops_from_gpu_span": 2.80504280155642, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.016416 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.081728, + "synchronized_e2e_ms": 0.10288 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.016448 + }, + "host_enqueue_ms": { + "median": 0.046208 + }, + "inter_kernel_gap_ms": { + "median": 3.2e-05 + }, + "kernel_sum_ms": { + "median": 0.016544 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 6232, + "submission_ms": { + "median": 0.046208 + }, + "synchronized_e2e_ms": { + "median": 0.06544 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.68314, + "after_init_synchronized_e2e_ms_per_call": 3.710884, + "including_init_host_enqueue_ms_per_call": 38.557224, + "including_init_synchronized_e2e_ms_per_call": 38.65396, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.4099012, + "after_init_synchronized_e2e_ms_per_call": 0.42998440000000004, + "including_init_host_enqueue_ms_per_call": 3.8973096, + "including_init_synchronized_e2e_ms_per_call": 3.924292, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.08257732000000001, + "after_init_synchronized_e2e_ms_per_call": 0.10189444, + "including_init_host_enqueue_ms_per_call": 0.43131816, + "including_init_synchronized_e2e_ms_per_call": 0.4513252, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.049844932, + "after_init_synchronized_e2e_ms_per_call": 0.069085444, + "including_init_host_enqueue_ms_per_call": 0.084719016, + "including_init_synchronized_e2e_ms_per_call": 0.10402852, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.68314, + "synchronized_e2e_ms": 3.710884, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.017152, + "median": 0.016448, + "min": 0.015712, + "p90": 0.016768, + "sample_count": 6232 + }, + "host_enqueue_ms": { + "max": 113.472758, + "median": 0.046208, + "min": 0.036384, + "p90": 0.0789600000000001, + "sample_count": 6232 + }, + "sample_count": 6232, + "synchronized_e2e_ms": { + "max": 113.786774, + "median": 0.06544, + "min": 0.056416, + "p90": 0.10143040000000006, + "sample_count": 6232 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.077696, + "submission_ms": 0.077696, + "synchronized_e2e_ms": 0.098208 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.05392, + "submission_ms": 0.05392, + "synchronized_e2e_ms": 0.071136 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.24736, + "submission_ms": 0.24736, + "synchronized_e2e_ms": 0.270944 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.196713, + "submission_ms": 8.196713, + "synchronized_e2e_ms": 8.224233 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.08736, + "submission_ms": 0.08736, + "synchronized_e2e_ms": 0.106272 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.050624, + "submission_ms": 0.050624, + "synchronized_e2e_ms": 0.067136 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.44093, + "submission_ms": 1.44093, + "synchronized_e2e_ms": 1.467522 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.200961, + "submission_ms": 1.200961, + "synchronized_e2e_ms": 1.223745 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.081728, + "submission_ms": 0.081728, + "synchronized_e2e_ms": 0.10288 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.68314, + "submission_ms": 3.68314, + "synchronized_e2e_ms": 3.710884 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.627167, + "evolution_kernel_ms": 0.172703, + "evolution_speedup": 3.6315, + "evolution_tflops": 0.2671, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d352_b1_n256_k256_d352", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 28784, + "measurement_schedule_sha256": "dd5745d2ec4ccbf9a67494ed1a260eed71396b3ef10bfa22a2afd6d01708ce7e", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 8160, + "public_pair_count": 6232, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 8160, + "baseline_public_raw": 6232, + "candidate_precomputed": 8160, + "candidate_public_raw": 6232 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 5758 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:post_d895_d352_b1_n256_k256_d352", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.6308943089430894, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.5760467603911983, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.2162463175890164, + "including_init_synchronized_e2e_speedup": 12.735195229673753, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.2655290517516447, + "including_init_synchronized_e2e_speedup": 12.582726145251169, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.4450113764794237, + "including_init_synchronized_e2e_speedup": 11.27689746772394, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.5567202911802958, + "including_init_synchronized_e2e_speedup": 6.35088964545492, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 3.348249027237354, + "hot_synchronized_e2e_speedup": 2.5760467603911983, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 235204, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d352_b1_n256_k256_d352", + "source": "high_small_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 352, + "K": 8192, + "N": 512, + "baseline_07cf_adapter_bench_iters": 2340, + "baseline_07cf_adapter_gpu_span_ms": 0.259361, + "baseline_07cf_adapter_host_enqueue_ms": 0.17292849999999999, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.038464, + "baseline_07cf_adapter_kernel_sum_ms": 0.22096, + "baseline_07cf_adapter_submission_ms": 0.17292849999999999, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.376464, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.22096 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.221888, + "synchronized_e2e_ms": 0.41584 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.259361 + }, + "host_enqueue_ms": { + "median": 0.17292849999999999 + }, + "inter_kernel_gap_ms": { + "median": 0.038464 + }, + "kernel_sum_ms": { + "median": 0.22096 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2340, + "submission_ms": { + "median": 0.17292849999999999 + }, + "synchronized_e2e_ms": { + "median": 0.376464 + } + }, + "baseline_07cf_precomputed_bench_iters": 3234, + "baseline_07cf_precomputed_gpu_span_ms": 0.234305, + "baseline_07cf_precomputed_host_enqueue_ms": 0.053952, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.234305, + "baseline_07cf_precomputed_submission_ms": 0.053952, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.2953765, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.234305 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.061825, + "synchronized_e2e_ms": 0.251617 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.234305 + }, + "host_enqueue_ms": { + "median": 0.053952 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.234305 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3234, + "submission_ms": { + "median": 0.053952 + }, + "synchronized_e2e_ms": { + "median": 0.2953765 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.1069375386782185, + "submission": 3.205228721826809, + "synchronized_e2e": 1.2745225161785045 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.658504, + "after_init_synchronized_e2e_ms_per_call": 7.839144, + "including_init_host_enqueue_ms_per_call": 43.209772, + "including_init_synchronized_e2e_ms_per_call": 43.470380999999996, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9214860500000001, + "after_init_synchronized_e2e_ms_per_call": 1.122732, + "including_init_host_enqueue_ms_per_call": 4.47661285, + "including_init_synchronized_e2e_ms_per_call": 4.685855699999999, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.247784255, + "after_init_synchronized_e2e_ms_per_call": 0.4510908, + "including_init_host_enqueue_ms_per_call": 0.603296935, + "including_init_synchronized_e2e_ms_per_call": 0.8074031700000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.18041407549999997, + "after_init_synchronized_e2e_ms_per_call": 0.38392667999999996, + "including_init_host_enqueue_ms_per_call": 0.21596534349999996, + "including_init_synchronized_e2e_ms_per_call": 0.419557917, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.658504, + "synchronized_e2e_ms": 7.839144, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.634112, + "median": 0.259361, + "min": 0.248161, + "p90": 0.28224319999999997, + "sample_count": 2340 + }, + "host_enqueue_ms": { + "max": 1.165825, + "median": 0.17292849999999999, + "min": 0.14304, + "p90": 0.21547929999999998, + "sample_count": 2340 + }, + "sample_count": 2340, + "synchronized_e2e_ms": { + "max": 72.068843, + "median": 0.376464, + "min": 0.349568, + "p90": 0.41454729999999995, + "sample_count": 2340 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3234, + "candidate_precomputed_gpu_span_ms": 0.030848, + "candidate_precomputed_host_enqueue_ms": 0.061728, + "candidate_precomputed_inter_kernel_gap_ms": 0.002112, + "candidate_precomputed_kernel_sum_ms": 0.028768, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.061728, + "candidate_precomputed_synchronized_e2e_ms": 0.075968, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.028768 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.062528, + "synchronized_e2e_ms": 0.08288 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.030848 + }, + "host_enqueue_ms": { + "median": 0.061728 + }, + "inter_kernel_gap_ms": { + "median": 0.002112 + }, + "kernel_sum_ms": { + "median": 0.028768 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3234, + "submission_ms": { + "median": 0.061728 + }, + "synchronized_e2e_ms": { + "median": 0.075968 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb93390bc0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb93393a40" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.383817427385892, + "submission": 1.029808190772421, + "synchronized_e2e": 1.4218618365627633 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2340, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.042688, + "candidate_public_raw_host_enqueue_ms": 0.063568, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.0428805, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.063568, + "candidate_public_raw_synchronized_e2e_ms": 0.108016, + "candidate_public_raw_tflops_from_gpu_span": 69.17143028485758, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.042688 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.259328, + "synchronized_e2e_ms": 0.288032 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.042688 + }, + "host_enqueue_ms": { + "median": 0.063568 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.0428805 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2340, + "submission_ms": { + "median": 0.063568 + }, + "synchronized_e2e_ms": { + "median": 0.108016 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.886467, + "after_init_synchronized_e2e_ms_per_call": 2.915524, + "including_init_host_enqueue_ms_per_call": 38.724487, + "including_init_synchronized_e2e_ms_per_call": 487.015737, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.34585790000000005, + "after_init_synchronized_e2e_ms_per_call": 0.3887668, + "including_init_host_enqueue_ms_per_call": 3.9296599, + "including_init_synchronized_e2e_ms_per_call": 48.7987881, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.09179699, + "after_init_synchronized_e2e_ms_per_call": 0.13609108, + "including_init_host_enqueue_ms_per_call": 0.45017719, + "including_init_synchronized_e2e_ms_per_call": 4.97709321, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.066390899, + "after_init_synchronized_e2e_ms_per_call": 0.110823508, + "including_init_host_enqueue_ms_per_call": 0.102228919, + "including_init_synchronized_e2e_ms_per_call": 0.594923721, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.886467, + "synchronized_e2e_ms": 2.915524, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.044064, + "median": 0.042688, + "min": 0.041824, + "p90": 0.043135999999999994, + "sample_count": 2340 + }, + "host_enqueue_ms": { + "max": 86.46793, + "median": 0.063568, + "min": 0.05088, + "p90": 0.08269119999999999, + "sample_count": 2340 + }, + "sample_count": 2340, + "synchronized_e2e_ms": { + "max": 96.748036, + "median": 0.108016, + "min": 0.095648, + "p90": 0.1243296, + "sample_count": 2340 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.086688, + "submission_ms": 0.086688, + "synchronized_e2e_ms": 0.278144 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.061825, + "submission_ms": 0.061825, + "synchronized_e2e_ms": 0.251617 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.221888, + "submission_ms": 0.221888, + "synchronized_e2e_ms": 0.41584 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.658504, + "submission_ms": 7.658504, + "synchronized_e2e_ms": 7.839144 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.090176, + "submission_ms": 0.090176, + "synchronized_e2e_ms": 0.109248 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.062528, + "submission_ms": 0.062528, + "synchronized_e2e_ms": 0.08288 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.32237, + "submission_ms": 1.32237, + "synchronized_e2e_ms": 1.348962 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.170433, + "submission_ms": 1.170433, + "synchronized_e2e_ms": 1.195169 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.259328, + "submission_ms": 0.259328, + "synchronized_e2e_ms": 0.288032 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.886467, + "submission_ms": 2.886467, + "synchronized_e2e_ms": 2.915524 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 47.910581, + "evolution_kernel_ms": 0.272767, + "evolution_speedup": 175.6465, + "evolution_tflops": 10.8253, + "expected_route": "d352_exactd_splitk_c95c_v2", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d352_b1_n512_k8192_d352", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 11148, + "measurement_schedule_sha256": "0b69c9f570d96cb3b2fea4d2d7f83840e4afe91e94cdabf6a207774eb593cbad", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3234, + "public_pair_count": 2340, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3234, + "baseline_public_raw": 2340, + "candidate_precomputed": 3234, + "candidate_public_raw": 2340 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2230 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:post_d895_d352_b1_n512_k8192_d352", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 7.595468101659751, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 3.4852614427492226, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.6887598935903116, + "including_init_synchronized_e2e_speedup": 0.08925867830837671, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.8879317884140314, + "including_init_synchronized_e2e_speedup": 0.09602401785875496, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 3.3146242942594033, + "including_init_synchronized_e2e_speedup": 0.1622238394848145, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 3.464307229834305, + "including_init_synchronized_e2e_speedup": 0.7052297667586195, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 6.075735569715143, + "hot_synchronized_e2e_speedup": 3.4852614427492226, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 235203, + "selected_route": "d352_exactd_splitk_c95c_v2", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d352_b1_n512_k8192_d352", + "source": "high_k_low_n", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 352, + "K": 1024, + "N": 2048, + "baseline_07cf_adapter_bench_iters": 2820, + "baseline_07cf_adapter_gpu_span_ms": 0.074336, + "baseline_07cf_adapter_host_enqueue_ms": 0.144385, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.031936, + "baseline_07cf_adapter_kernel_sum_ms": 0.042432, + "baseline_07cf_adapter_submission_ms": 0.144385, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.172544, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.042432 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.225313, + "synchronized_e2e_ms": 0.246561 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.074336 + }, + "host_enqueue_ms": { + "median": 0.144385 + }, + "inter_kernel_gap_ms": { + "median": 0.031936 + }, + "kernel_sum_ms": { + "median": 0.042432 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2820, + "submission_ms": { + "median": 0.144385 + }, + "synchronized_e2e_ms": { + "median": 0.172544 + } + }, + "baseline_07cf_precomputed_bench_iters": 3357, + "baseline_07cf_precomputed_gpu_span_ms": 0.04128, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042816, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.04128, + "baseline_07cf_precomputed_submission_ms": 0.042816, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.089728, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.04128 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.096704, + "synchronized_e2e_ms": 0.127392 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.04128 + }, + "host_enqueue_ms": { + "median": 0.042816 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.04128 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3357, + "submission_ms": { + "median": 0.042816 + }, + "synchronized_e2e_ms": { + "median": 0.089728 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.8007751937984497, + "submission": 3.372220665171899, + "synchronized_e2e": 1.9229671897289586 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 6.927783, + "after_init_synchronized_e2e_ms_per_call": 6.957031, + "including_init_host_enqueue_ms_per_call": 43.862221, + "including_init_synchronized_e2e_ms_per_call": 462.06835, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8227247999999999, + "after_init_synchronized_e2e_ms_per_call": 0.8509926999999999, + "including_init_host_enqueue_ms_per_call": 4.516168599999999, + "including_init_synchronized_e2e_ms_per_call": 46.3621246, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.21221898000000003, + "after_init_synchronized_e2e_ms_per_call": 0.24038887000000003, + "including_init_host_enqueue_ms_per_call": 0.58156336, + "including_init_synchronized_e2e_ms_per_call": 4.79150206, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15116839800000004, + "after_init_synchronized_e2e_ms_per_call": 0.179328487, + "including_init_host_enqueue_ms_per_call": 0.18810283600000002, + "including_init_synchronized_e2e_ms_per_call": 0.6344398059999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 6.927783, + "synchronized_e2e_ms": 6.957031, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.483777, + "median": 0.074336, + "min": 0.065888, + "p90": 0.09890329999999997, + "sample_count": 2820 + }, + "host_enqueue_ms": { + "max": 0.875905, + "median": 0.144385, + "min": 0.123136, + "p90": 0.2012833, + "sample_count": 2820 + }, + "sample_count": 2820, + "synchronized_e2e_ms": { + "max": 45.106447, + "median": 0.172544, + "min": 0.153504, + "p90": 0.2271848999999999, + "sample_count": 2820 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 3357, + "candidate_precomputed_gpu_span_ms": 0.029792, + "candidate_precomputed_host_enqueue_ms": 0.050816, + "candidate_precomputed_inter_kernel_gap_ms": 0.002272, + "candidate_precomputed_kernel_sum_ms": 0.02752, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.050816, + "candidate_precomputed_synchronized_e2e_ms": 0.064416, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.02752 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.040128, + "synchronized_e2e_ms": 0.055456 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.029792 + }, + "host_enqueue_ms": { + "median": 0.050816 + }, + "inter_kernel_gap_ms": { + "median": 0.002272 + }, + "kernel_sum_ms": { + "median": 0.02752 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 3357, + "submission_ms": { + "median": 0.050816 + }, + "synchronized_e2e_ms": { + "median": 0.064416 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc03ca20c0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01adf200" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.1911922663802363, + "submission": 0.922544080604534, + "synchronized_e2e": 1.2985593641331346 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2820, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.035488, + "candidate_public_raw_host_enqueue_ms": 0.04688, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.03552, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.04688, + "candidate_public_raw_synchronized_e2e_ms": 0.083648, + "candidate_public_raw_tflops_from_gpu_span": 83.20530928764653, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.035424 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.244544, + "synchronized_e2e_ms": 0.267936 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.035488 + }, + "host_enqueue_ms": { + "median": 0.04688 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.03552 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2820, + "submission_ms": { + "median": 0.04688 + }, + "synchronized_e2e_ms": { + "median": 0.083648 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 68.441095, + "after_init_synchronized_e2e_ms_per_call": 68.472839, + "including_init_host_enqueue_ms_per_call": 105.79911799999999, + "including_init_synchronized_e2e_ms_per_call": 105.90846199999999, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 6.8863015, + "after_init_synchronized_e2e_ms_per_call": 6.922567099999999, + "including_init_host_enqueue_ms_per_call": 10.6221038, + "including_init_synchronized_e2e_ms_per_call": 10.666129399999999, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.7308221500000001, + "after_init_synchronized_e2e_ms_per_call": 0.76753991, + "including_init_host_enqueue_ms_per_call": 1.10440238, + "including_init_synchronized_e2e_ms_per_call": 1.1418961399999998, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.115274215, + "after_init_synchronized_e2e_ms_per_call": 0.15203719100000002, + "including_init_host_enqueue_ms_per_call": 0.152632238, + "including_init_synchronized_e2e_ms_per_call": 0.18947281399999996, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 68.441095, + "synchronized_e2e_ms": 68.472839, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.036768, + "median": 0.035488, + "min": 0.034464, + "p90": 0.035936, + "sample_count": 2820 + }, + "host_enqueue_ms": { + "max": 0.424288, + "median": 0.04688, + "min": 0.036896, + "p90": 0.07158719999999999, + "sample_count": 2820 + }, + "sample_count": 2820, + "synchronized_e2e_ms": { + "max": 0.531201, + "median": 0.083648, + "min": 0.074784, + "p90": 0.10584959999999997, + "sample_count": 2820 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.06544, + "submission_ms": 0.06544, + "synchronized_e2e_ms": 0.095649 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.096704, + "submission_ms": 0.096704, + "synchronized_e2e_ms": 0.127392 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.225313, + "submission_ms": 0.225313, + "synchronized_e2e_ms": 0.246561 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 6.927783, + "submission_ms": 6.927783, + "synchronized_e2e_ms": 6.957031 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.064288, + "submission_ms": 0.064288, + "synchronized_e2e_ms": 0.082336 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.040128, + "submission_ms": 0.040128, + "synchronized_e2e_ms": 0.055456 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.242049, + "submission_ms": 1.242049, + "synchronized_e2e_ms": 1.263169 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.235777, + "submission_ms": 1.235777, + "synchronized_e2e_ms": 1.257473 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.244544, + "submission_ms": 0.244544, + "synchronized_e2e_ms": 0.267936 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 68.441095, + "submission_ms": 68.441095, + "synchronized_e2e_ms": 68.472839 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.604831, + "evolution_kernel_ms": 0.189184, + "evolution_speedup": 3.1971, + "evolution_tflops": 15.608, + "expected_route": "d352_exactd_splitk_c95c_v2", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d352_b2_n2048_k1024_d352", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 12354, + "measurement_schedule_sha256": "ee98120106d624954d94513534fc5dd7d32f95fb56087b37fcd330a6f67fa6a7", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 3357, + "public_pair_count": 2820, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 3357, + "baseline_public_raw": 2820, + "candidate_precomputed": 3357, + "candidate_public_raw": 2820 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2472 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:post_d895_d352_b2_n2048_k1024_d352", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.3856068743286787, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.0627390971690898, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.10160278296625032, + "including_init_synchronized_e2e_speedup": 4.3629030322430715, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.12293022049580422, + "including_init_synchronized_e2e_speedup": 4.3466681174897435, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.31319396798532606, + "including_init_synchronized_e2e_speedup": 4.196092702441398, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.1795040793670017, + "including_init_synchronized_e2e_speedup": 3.3484476881205767, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.0946798917944096, + "hot_synchronized_e2e_speedup": 2.0627390971690898, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 235201, + "selected_route": "d352_exactd_splitk_c95c_v2", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d352_b2_n2048_k1024_d352", + "source": "new_high_gap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 352, + "K": 1024, + "N": 8192, + "baseline_07cf_adapter_bench_iters": 880, + "baseline_07cf_adapter_gpu_span_ms": 0.1729125, + "baseline_07cf_adapter_host_enqueue_ms": 0.147344, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.002176, + "baseline_07cf_adapter_kernel_sum_ms": 0.1705605, + "baseline_07cf_adapter_submission_ms": 0.147344, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.2729445, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.1705605 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.188768, + "synchronized_e2e_ms": 0.297376 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.1729125 + }, + "host_enqueue_ms": { + "median": 0.147344 + }, + "inter_kernel_gap_ms": { + "median": 0.002176 + }, + "kernel_sum_ms": { + "median": 0.1705605 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 880, + "submission_ms": { + "median": 0.147344 + }, + "synchronized_e2e_ms": { + "median": 0.2729445 + } + }, + "baseline_07cf_precomputed_bench_iters": 1580, + "baseline_07cf_precomputed_gpu_span_ms": 0.127456, + "baseline_07cf_precomputed_host_enqueue_ms": 0.041728, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.127456, + "baseline_07cf_precomputed_submission_ms": 0.041728, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.17568, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.127456 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.044896, + "synchronized_e2e_ms": 0.169952 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.127456 + }, + "host_enqueue_ms": { + "median": 0.041728 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.127456 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1580, + "submission_ms": { + "median": 0.041728 + }, + "synchronized_e2e_ms": { + "median": 0.17568 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.3566446459954806, + "submission": 3.5310582822085887, + "synchronized_e2e": 1.5536458333333332 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.163496, + "after_init_synchronized_e2e_ms_per_call": 8.263849, + "including_init_host_enqueue_ms_per_call": 42.391403000000004, + "including_init_synchronized_e2e_ms_per_call": 42.609292, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9489592, + "after_init_synchronized_e2e_ms_per_call": 1.0720349500000002, + "including_init_host_enqueue_ms_per_call": 4.3717499, + "including_init_synchronized_e2e_ms_per_call": 4.50657925, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22750552, + "after_init_synchronized_e2e_ms_per_call": 0.352853545, + "including_init_host_enqueue_ms_per_call": 0.56978459, + "including_init_synchronized_e2e_ms_per_call": 0.696307975, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.155360152, + "after_init_synchronized_e2e_ms_per_call": 0.28093540449999993, + "including_init_host_enqueue_ms_per_call": 0.18958805899999998, + "including_init_synchronized_e2e_ms_per_call": 0.31528084749999996, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.163496, + "synchronized_e2e_ms": 8.263849, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.423168, + "median": 0.1729125, + "min": 0.165376, + "p90": 0.1756512, + "sample_count": 880 + }, + "host_enqueue_ms": { + "max": 0.440704, + "median": 0.147344, + "min": 0.128384, + "p90": 0.1785728, + "sample_count": 880 + }, + "sample_count": 880, + "synchronized_e2e_ms": { + "max": 0.53936, + "median": 0.2729445, + "min": 0.254464, + "p90": 0.295392, + "sample_count": 880 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 1580, + "candidate_precomputed_gpu_span_ms": 0.063616, + "candidate_precomputed_host_enqueue_ms": 0.052064, + "candidate_precomputed_inter_kernel_gap_ms": 0.002144, + "candidate_precomputed_kernel_sum_ms": 0.06144, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.052064, + "candidate_precomputed_synchronized_e2e_ms": 0.097536, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.06144 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.03904, + "synchronized_e2e_ms": 0.089568 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.063616 + }, + "host_enqueue_ms": { + "median": 0.052064 + }, + "inter_kernel_gap_ms": { + "median": 0.002144 + }, + "kernel_sum_ms": { + "median": 0.06144 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1580, + "submission_ms": { + "median": 0.052064 + }, + "synchronized_e2e_ms": { + "median": 0.097536 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc03730b90", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc03732c30" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.7872233400402413, + "submission": 0.8500307314074984, + "synchronized_e2e": 1.6443569553805775 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 880, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.113696, + "candidate_public_raw_host_enqueue_ms": 0.044256, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.113888, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.044256, + "candidate_public_raw_synchronized_e2e_ms": 0.160384, + "candidate_public_raw_tflops_from_gpu_span": 207.76738080495355, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.113696 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.049568, + "synchronized_e2e_ms": 0.141824 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.113696 + }, + "host_enqueue_ms": { + "median": 0.044256 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.113888 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 880, + "submission_ms": { + "median": 0.044256 + }, + "synchronized_e2e_ms": { + "median": 0.160384 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.002114, + "after_init_synchronized_e2e_ms_per_call": 2.088322, + "including_init_host_enqueue_ms_per_call": 36.579654, + "including_init_synchronized_e2e_ms_per_call": 450.176114, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.24004180000000003, + "after_init_synchronized_e2e_ms_per_call": 0.3531778, + "including_init_host_enqueue_ms_per_call": 3.6977958, + "including_init_synchronized_e2e_ms_per_call": 45.161957, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06383457999999999, + "after_init_synchronized_e2e_ms_per_call": 0.17966338, + "including_init_host_enqueue_ms_per_call": 0.40960998, + "including_init_synchronized_e2e_ms_per_call": 4.6605413, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.046213858, + "after_init_synchronized_e2e_ms_per_call": 0.162311938, + "including_init_host_enqueue_ms_per_call": 0.08079139799999999, + "including_init_synchronized_e2e_ms_per_call": 0.61039973, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.002114, + "synchronized_e2e_ms": 2.088322, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.116064, + "median": 0.113696, + "min": 0.111552, + "p90": 0.114144, + "sample_count": 880 + }, + "host_enqueue_ms": { + "max": 0.088128, + "median": 0.044256, + "min": 0.036224, + "p90": 0.052128, + "sample_count": 880 + }, + "sample_count": 880, + "synchronized_e2e_ms": { + "max": 0.195968, + "median": 0.160384, + "min": 0.15344, + "p90": 0.167296, + "sample_count": 880 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.08128, + "submission_ms": 0.08128, + "synchronized_e2e_ms": 0.207328 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.044896, + "submission_ms": 0.044896, + "synchronized_e2e_ms": 0.169952 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.188768, + "submission_ms": 0.188768, + "synchronized_e2e_ms": 0.297376 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.163496, + "submission_ms": 8.163496, + "synchronized_e2e_ms": 8.263849 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.059808, + "submission_ms": 0.059808, + "synchronized_e2e_ms": 0.099936 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.03904, + "submission_ms": 0.03904, + "synchronized_e2e_ms": 0.089568 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.160513, + "submission_ms": 1.160513, + "synchronized_e2e_ms": 1.185377 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.085378, + "submission_ms": 1.085378, + "synchronized_e2e_ms": 1.10605 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.049568, + "submission_ms": 0.049568, + "synchronized_e2e_ms": 0.141824 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.002114, + "submission_ms": 2.002114, + "synchronized_e2e_ms": 2.088322 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.651839, + "evolution_kernel_ms": 0.220992, + "evolution_speedup": 2.9496, + "evolution_tflops": 106.8922, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d352_b4_n8192_k1024_d352", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 4920, + "measurement_schedule_sha256": "babb2d8de9e4e723afa71da7ce6ce8aeb514bbf399fba619803a3883b0d3e18e", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 1580, + "public_pair_count": 880, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 1580, + "baseline_public_raw": 880, + "candidate_precomputed": 1580, + "candidate_public_raw": 880 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 984 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:post_d895_d352_b4_n8192_k1024_d352", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 2.0035211267605635, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.7018187599760572, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.957171834611713, + "including_init_synchronized_e2e_speedup": 0.09465027280412307, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 3.035397326785546, + "including_init_synchronized_e2e_speedup": 0.09978706746476906, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 1.9639703149300651, + "including_init_synchronized_e2e_speedup": 0.14940495753143523, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.7308363633733457, + "including_init_synchronized_e2e_speedup": 0.5165153783734471, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.5208318674359695, + "hot_synchronized_e2e_speedup": 1.7018187599760572, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 235202, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d352_b4_n8192_k1024_d352", + "source": "new_high_gap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 416, + "K": 256, + "N": 256, + "baseline_07cf_adapter_bench_iters": 5825, + "baseline_07cf_adapter_gpu_span_ms": 0.05664, + "baseline_07cf_adapter_host_enqueue_ms": 0.152129, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.043136, + "baseline_07cf_adapter_kernel_sum_ms": 0.013504, + "baseline_07cf_adapter_submission_ms": 0.152129, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.173088, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.013504 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.184928, + "synchronized_e2e_ms": 0.20272 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.05664 + }, + "host_enqueue_ms": { + "median": 0.152129 + }, + "inter_kernel_gap_ms": { + "median": 0.043136 + }, + "kernel_sum_ms": { + "median": 0.013504 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 5825, + "submission_ms": { + "median": 0.152129 + }, + "synchronized_e2e_ms": { + "median": 0.173088 + } + }, + "baseline_07cf_precomputed_bench_iters": 8139, + "baseline_07cf_precomputed_gpu_span_ms": 0.012448, + "baseline_07cf_precomputed_host_enqueue_ms": 0.044512, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.012448, + "baseline_07cf_precomputed_submission_ms": 0.044512, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.064353, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.012448 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.03984, + "synchronized_e2e_ms": 0.056096 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.012448 + }, + "host_enqueue_ms": { + "median": 0.044512 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.012448 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 8139, + "submission_ms": { + "median": 0.044512 + }, + "synchronized_e2e_ms": { + "median": 0.064353 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 4.5501285347043705, + "submission": 3.4177075844716027, + "synchronized_e2e": 2.689664817491026 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.424328, + "after_init_synchronized_e2e_ms_per_call": 8.453128, + "including_init_host_enqueue_ms_per_call": 42.873772, + "including_init_synchronized_e2e_ms_per_call": 492.494622, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9793489, + "after_init_synchronized_e2e_ms_per_call": 1.0010919999999999, + "including_init_host_enqueue_ms_per_call": 4.4242933, + "including_init_synchronized_e2e_ms_per_call": 49.4052414, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23485098999999998, + "after_init_synchronized_e2e_ms_per_call": 0.25588839999999996, + "including_init_host_enqueue_ms_per_call": 0.57934543, + "including_init_synchronized_e2e_ms_per_call": 5.09630334, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.160401199, + "after_init_synchronized_e2e_ms_per_call": 0.18136803999999998, + "including_init_host_enqueue_ms_per_call": 0.194850643, + "including_init_synchronized_e2e_ms_per_call": 0.665409534, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.424328, + "synchronized_e2e_ms": 8.453128, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.473377, + "median": 0.05664, + "min": 0.044416, + "p90": 0.08668160000000004, + "sample_count": 5825 + }, + "host_enqueue_ms": { + "max": 101.898441, + "median": 0.152129, + "min": 0.126528, + "p90": 0.25947580000000015, + "sample_count": 5825 + }, + "sample_count": 5825, + "synchronized_e2e_ms": { + "max": 102.155657, + "median": 0.173088, + "min": 0.144448, + "p90": 0.28986880000000004, + "sample_count": 5825 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 8139, + "candidate_precomputed_gpu_span_ms": 0.020672, + "candidate_precomputed_host_enqueue_ms": 0.057248, + "candidate_precomputed_inter_kernel_gap_ms": 0.00688, + "candidate_precomputed_kernel_sum_ms": 0.013824, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.057248, + "candidate_precomputed_synchronized_e2e_ms": 0.070368, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.013824 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.04, + "synchronized_e2e_ms": 0.053504 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.020672 + }, + "host_enqueue_ms": { + "median": 0.057248 + }, + "inter_kernel_gap_ms": { + "median": 0.00688 + }, + "kernel_sum_ms": { + "median": 0.013824 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 8139, + "submission_ms": { + "median": 0.057248 + }, + "synchronized_e2e_ms": { + "median": 0.070368 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb731860c0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb73185670" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 0.848297213622291, + "submission": 0.8261598658468418, + "synchronized_e2e": 0.9495225102319236 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 5825, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.017536, + "candidate_public_raw_host_enqueue_ms": 0.047296, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.017472, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.047296, + "candidate_public_raw_synchronized_e2e_ms": 0.066816, + "candidate_public_raw_tflops_from_gpu_span": 3.109372262773723, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.017472 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.04992, + "synchronized_e2e_ms": 0.065248 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.017536 + }, + "host_enqueue_ms": { + "median": 0.047296 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.017472 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 5825, + "submission_ms": { + "median": 0.047296 + }, + "synchronized_e2e_ms": { + "median": 0.066816 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.06378, + "after_init_synchronized_e2e_ms_per_call": 3.088356, + "including_init_host_enqueue_ms_per_call": 37.937864, + "including_init_synchronized_e2e_ms_per_call": 38.031431999999995, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.3489444, + "after_init_synchronized_e2e_ms_per_call": 0.36897, + "including_init_host_enqueue_ms_per_call": 3.8363527999999993, + "including_init_synchronized_e2e_ms_per_call": 3.863277599999999, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07746084, + "after_init_synchronized_e2e_ms_per_call": 0.09703140000000002, + "including_init_host_enqueue_ms_per_call": 0.42620167999999997, + "including_init_synchronized_e2e_ms_per_call": 0.44646215999999994, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.050312484, + "after_init_synchronized_e2e_ms_per_call": 0.06983754, + "including_init_host_enqueue_ms_per_call": 0.08518656799999999, + "including_init_synchronized_e2e_ms_per_call": 0.104780616, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.06378, + "synchronized_e2e_ms": 3.088356, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.018208, + "median": 0.017536, + "min": 0.016736, + "p90": 0.01808, + "sample_count": 5825 + }, + "host_enqueue_ms": { + "max": 98.829127, + "median": 0.047296, + "min": 0.036032, + "p90": 0.08236160000000003, + "sample_count": 5825 + }, + "sample_count": 5825, + "synchronized_e2e_ms": { + "max": 99.165479, + "median": 0.066816, + "min": 0.056928, + "p90": 0.10590780000000002, + "sample_count": 5825 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.06544, + "submission_ms": 0.06544, + "synchronized_e2e_ms": 0.082304 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.03984, + "submission_ms": 0.03984, + "synchronized_e2e_ms": 0.056096 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.184928, + "submission_ms": 0.184928, + "synchronized_e2e_ms": 0.20272 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.424328, + "submission_ms": 8.424328, + "synchronized_e2e_ms": 8.453128 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.080768, + "submission_ms": 0.080768, + "synchronized_e2e_ms": 0.096032 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.04, + "submission_ms": 0.04, + "synchronized_e2e_ms": 0.053504 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.239073, + "submission_ms": 1.239073, + "synchronized_e2e_ms": 1.265633 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.041633, + "submission_ms": 1.041633, + "synchronized_e2e_ms": 1.063233 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.04992, + "submission_ms": 0.04992, + "synchronized_e2e_ms": 0.065248 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.06378, + "submission_ms": 3.06378, + "synchronized_e2e_ms": 3.088356 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.647807, + "evolution_kernel_ms": 0.175137, + "evolution_speedup": 3.6989, + "evolution_tflops": 0.3113, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d416_b1_n256_k256_d416", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 27928, + "measurement_schedule_sha256": "e84881fc74655c01e02c8992f0161edc01cbd8a66916d678e99f7263f2df3d04", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 8139, + "public_pair_count": 5825, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 8139, + "baseline_public_raw": 5825, + "candidate_precomputed": 8139, + "candidate_public_raw": 5825 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 5586 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:post_d895_d416_b1_n256_k256_d416", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.6021671826625388, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.59051724137931, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.737096371014222, + "including_init_synchronized_e2e_speedup": 12.949673364915633, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.7132070358023683, + "including_init_synchronized_e2e_speedup": 12.788426438731717, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.6371710600898255, + "including_init_synchronized_e2e_speedup": 11.414860645748794, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.5969992642925277, + "including_init_synchronized_e2e_speedup": 6.350502215028017, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 3.2299270072992705, + "hot_synchronized_e2e_speedup": 2.59051724137931, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 241604, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d416_b1_n256_k256_d416", + "source": "high_small_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 416, + "K": 8192, + "N": 512, + "baseline_07cf_adapter_bench_iters": 2382, + "baseline_07cf_adapter_gpu_span_ms": 0.2687525, + "baseline_07cf_adapter_host_enqueue_ms": 0.17449599999999998, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.039584, + "baseline_07cf_adapter_kernel_sum_ms": 0.229089, + "baseline_07cf_adapter_submission_ms": 0.17449599999999998, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.386208, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.229089 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.249728, + "synchronized_e2e_ms": 0.449984 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.2687525 + }, + "host_enqueue_ms": { + "median": 0.17449599999999998 + }, + "inter_kernel_gap_ms": { + "median": 0.039584 + }, + "kernel_sum_ms": { + "median": 0.229089 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2382, + "submission_ms": { + "median": 0.17449599999999998 + }, + "synchronized_e2e_ms": { + "median": 0.386208 + } + }, + "baseline_07cf_precomputed_bench_iters": 2000, + "baseline_07cf_precomputed_gpu_span_ms": 0.241216, + "baseline_07cf_precomputed_host_enqueue_ms": 0.055456, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.241216, + "baseline_07cf_precomputed_submission_ms": 0.055456, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.30368, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.241216 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.073632, + "synchronized_e2e_ms": 0.271808 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.241216 + }, + "host_enqueue_ms": { + "median": 0.055456 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.241216 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2000, + "submission_ms": { + "median": 0.055456 + }, + "synchronized_e2e_ms": { + "median": 0.30368 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.1141570210931282, + "submission": 3.1465666474321985, + "synchronized_e2e": 1.2717597471022128 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.632265, + "after_init_synchronized_e2e_ms_per_call": 8.81761, + "including_init_host_enqueue_ms_per_call": 44.183533, + "including_init_synchronized_e2e_ms_per_call": 44.448847, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 1.0202729, + "after_init_synchronized_e2e_ms_per_call": 1.2293482, + "including_init_host_enqueue_ms_per_call": 4.5753997, + "including_init_synchronized_e2e_ms_per_call": 4.792471900000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.25907368999999997, + "after_init_synchronized_e2e_ms_per_call": 0.47052202, + "including_init_host_enqueue_ms_per_call": 0.6145863699999999, + "including_init_synchronized_e2e_ms_per_call": 0.8268343899999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.18295376899999996, + "after_init_synchronized_e2e_ms_per_call": 0.39463940200000003, + "including_init_host_enqueue_ms_per_call": 0.21850503699999996, + "including_init_synchronized_e2e_ms_per_call": 0.430270639, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.632265, + "synchronized_e2e_ms": 8.81761, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 10.085098, + "median": 0.2687525, + "min": 0.257472, + "p90": 0.28828580000000004, + "sample_count": 2382 + }, + "host_enqueue_ms": { + "max": 41.168299, + "median": 0.17449599999999998, + "min": 0.145216, + "p90": 0.2095328, + "sample_count": 2382 + }, + "sample_count": 2382, + "synchronized_e2e_ms": { + "max": 41.334539, + "median": 0.386208, + "min": 0.36064, + "p90": 0.41804250000000004, + "sample_count": 2382 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 4, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 2000, + "candidate_precomputed_gpu_span_ms": 0.035616, + "candidate_precomputed_host_enqueue_ms": 0.077072, + "candidate_precomputed_inter_kernel_gap_ms": 0.008384, + "candidate_precomputed_kernel_sum_ms": 0.027167999999999998, + "candidate_precomputed_launch_count": 3, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.077072, + "candidate_precomputed_synchronized_e2e_ms": 0.091264, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.027167999999999998 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.065184, + "synchronized_e2e_ms": 0.080896 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 3.0 + }, + "gpu_span_ms": { + "median": 0.035616 + }, + "host_enqueue_ms": { + "median": 0.077072 + }, + "inter_kernel_gap_ms": { + "median": 0.008384 + }, + "kernel_sum_ms": { + "median": 0.027167999999999998 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2000, + "submission_ms": { + "median": 0.077072 + }, + "synchronized_e2e_ms": { + "median": 0.091264 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01233f20", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc01232b70" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.1778975741239892, + "submission": 0.8374506954536018, + "synchronized_e2e": 1.1840813464235624 + }, + "candidate_public_raw_assignment_launch_count": 3, + "candidate_public_raw_bench_iters": 2382, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.041952, + "candidate_public_raw_host_enqueue_ms": 0.064544, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.041984, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.064544, + "candidate_public_raw_synchronized_e2e_ms": 0.108064, + "candidate_public_raw_tflops_from_gpu_span": 83.18223035850495, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.041856 + }, + "activity_count": { + "median": 4.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.07024, + "synchronized_e2e_ms": 0.100128 + }, + "correlated_kernel_activity_count": { + "median": 4.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.041952 + }, + "host_enqueue_ms": { + "median": 0.064544 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.041984 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2382, + "submission_ms": { + "median": 0.064544 + }, + "synchronized_e2e_ms": { + "median": 0.108064 + } + }, + "candidate_public_raw_total_launch_count": 4, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 73.52798, + "after_init_synchronized_e2e_ms_per_call": 73.5595, + "including_init_host_enqueue_ms_per_call": 109.366, + "including_init_synchronized_e2e_ms_per_call": 557.659713, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 7.4108876, + "after_init_synchronized_e2e_ms_per_call": 7.453207600000001, + "including_init_host_enqueue_ms_per_call": 10.9946896, + "including_init_synchronized_e2e_ms_per_call": 55.8632289, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.79917836, + "after_init_synchronized_e2e_ms_per_call": 0.84257836, + "including_init_host_enqueue_ms_per_call": 1.15755856, + "including_init_synchronized_e2e_ms_per_call": 5.683580490000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.13800743599999998, + "after_init_synchronized_e2e_ms_per_call": 0.181515436, + "including_init_host_enqueue_ms_per_call": 0.173845456, + "including_init_synchronized_e2e_ms_per_call": 0.6656156489999999, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 73.52798, + "synchronized_e2e_ms": 73.5595, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.043168, + "median": 0.041952, + "min": 0.041504, + "p90": 0.042272, + "sample_count": 2382 + }, + "host_enqueue_ms": { + "max": 0.767617, + "median": 0.064544, + "min": 0.044192, + "p90": 0.08185600000000001, + "sample_count": 2382 + }, + "sample_count": 2382, + "synchronized_e2e_ms": { + "max": 0.848513, + "median": 0.108064, + "min": 0.089376, + "p90": 0.12286730000000001, + "sample_count": 2382 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.081376, + "submission_ms": 0.081376, + "synchronized_e2e_ms": 0.276577 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.073632, + "submission_ms": 0.073632, + "synchronized_e2e_ms": 0.271808 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.249728, + "submission_ms": 0.249728, + "synchronized_e2e_ms": 0.449984 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.632265, + "submission_ms": 8.632265, + "synchronized_e2e_ms": 8.81761 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.113184, + "submission_ms": 0.113184, + "synchronized_e2e_ms": 0.133312 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.065184, + "submission_ms": 0.065184, + "synchronized_e2e_ms": 0.080896 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.383394, + "submission_ms": 1.383394, + "synchronized_e2e_ms": 1.409602 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.166401, + "submission_ms": 1.166401, + "synchronized_e2e_ms": 1.190689 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.07024, + "submission_ms": 0.07024, + "synchronized_e2e_ms": 0.100128 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 73.52798, + "submission_ms": 73.52798, + "synchronized_e2e_ms": 73.5595 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 40.967689, + "evolution_kernel_ms": 0.275103, + "evolution_speedup": 148.9176, + "evolution_tflops": 12.6849, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d416_b1_n512_k8192_d416", + "measurement_order": [ + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 8764, + "measurement_schedule_sha256": "abcbe97c0dde8b01931512a5d57302f96115ccfe2e9df08b7c82cbd7e6574344", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 2000, + "public_pair_count": 2382, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 2000, + "baseline_public_raw": 2382, + "candidate_precomputed": 2000, + "candidate_public_raw": 2382 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 1754 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:post_d895_d416_b1_n512_k8192_d416", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 6.772686433063791, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 3.5738821439147173, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.11987044501390032, + "including_init_synchronized_e2e_speedup": 0.07970603929927425, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.16494216530343256, + "including_init_synchronized_e2e_speedup": 0.08578938228899978, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.5584311707222104, + "including_init_synchronized_e2e_speedup": 0.14547773035937067, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.174136870651596, + "including_init_synchronized_e2e_speedup": 0.6464250647448345, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 6.406190408085431, + "hot_synchronized_e2e_speedup": 3.5738821439147173, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 241603, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d416_b1_n512_k8192_d416", + "source": "high_k_low_n", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 416, + "K": 1024, + "N": 2048, + "baseline_07cf_adapter_bench_iters": 2174, + "baseline_07cf_adapter_gpu_span_ms": 0.074688, + "baseline_07cf_adapter_host_enqueue_ms": 0.143616, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.031504500000000005, + "baseline_07cf_adapter_kernel_sum_ms": 0.043168, + "baseline_07cf_adapter_submission_ms": 0.143616, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.17201699999999998, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.043168 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.193792, + "synchronized_e2e_ms": 0.217248 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.074688 + }, + "host_enqueue_ms": { + "median": 0.143616 + }, + "inter_kernel_gap_ms": { + "median": 0.031504500000000005 + }, + "kernel_sum_ms": { + "median": 0.043168 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2174, + "submission_ms": { + "median": 0.143616 + }, + "synchronized_e2e_ms": { + "median": 0.17201699999999998 + } + }, + "baseline_07cf_precomputed_bench_iters": 2521, + "baseline_07cf_precomputed_gpu_span_ms": 0.042753, + "baseline_07cf_precomputed_host_enqueue_ms": 0.042432, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.042753, + "baseline_07cf_precomputed_submission_ms": 0.042432, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.09088, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.042753 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.044288, + "synchronized_e2e_ms": 0.076512 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.042753 + }, + "host_enqueue_ms": { + "median": 0.042432 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.042753 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2521, + "submission_ms": { + "median": 0.042432 + }, + "synchronized_e2e_ms": { + "median": 0.09088 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.7469651252543683, + "submission": 3.3846153846153846, + "synchronized_e2e": 1.8927926936619714 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.225448, + "after_init_synchronized_e2e_ms_per_call": 7.252648, + "including_init_host_enqueue_ms_per_call": 44.159886, + "including_init_synchronized_e2e_ms_per_call": 462.36396700000006, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8517992, + "after_init_synchronized_e2e_ms_per_call": 0.8800801, + "including_init_host_enqueue_ms_per_call": 4.545243, + "including_init_synchronized_e2e_ms_per_call": 46.39121200000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.21443432, + "after_init_synchronized_e2e_ms_per_call": 0.24282331, + "including_init_host_enqueue_ms_per_call": 0.5837787, + "including_init_synchronized_e2e_ms_per_call": 4.7939365, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.150697832, + "after_init_synchronized_e2e_ms_per_call": 0.17909763099999998, + "including_init_host_enqueue_ms_per_call": 0.18763227000000002, + "including_init_synchronized_e2e_ms_per_call": 0.6342089500000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.225448, + "synchronized_e2e_ms": 7.252648, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 5.287333, + "median": 0.074688, + "min": 0.065344, + "p90": 0.0939648, + "sample_count": 2174 + }, + "host_enqueue_ms": { + "max": 36.414886, + "median": 0.143616, + "min": 0.123488, + "p90": 0.17924800000000002, + "sample_count": 2174 + }, + "sample_count": 2174, + "synchronized_e2e_ms": { + "max": 44.645998, + "median": 0.17201699999999998, + "min": 0.152865, + "p90": 0.2066791, + "sample_count": 2174 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 2521, + "candidate_precomputed_gpu_span_ms": 0.037504, + "candidate_precomputed_host_enqueue_ms": 0.05344, + "candidate_precomputed_inter_kernel_gap_ms": 0.003616, + "candidate_precomputed_kernel_sum_ms": 0.033856, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.05344, + "candidate_precomputed_synchronized_e2e_ms": 0.071136, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.033856 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.042528, + "synchronized_e2e_ms": 0.065504 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.037504 + }, + "host_enqueue_ms": { + "median": 0.05344 + }, + "inter_kernel_gap_ms": { + "median": 0.003616 + }, + "kernel_sum_ms": { + "median": 0.033856 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2521, + "submission_ms": { + "median": 0.05344 + }, + "synchronized_e2e_ms": { + "median": 0.071136 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb6b3d1010", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb6b3d0920" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.226962457337884, + "submission": 0.8215568862275449, + "synchronized_e2e": 1.285657051282051 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2174, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.046016, + "candidate_public_raw_host_enqueue_ms": 0.043904, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.046208, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.043904, + "candidate_public_raw_synchronized_e2e_ms": 0.0914565, + "candidate_public_raw_tflops_from_gpu_span": 75.8358164116829, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.046016 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.069824, + "synchronized_e2e_ms": 0.122752 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.046016 + }, + "host_enqueue_ms": { + "median": 0.043904 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.046208 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2174, + "submission_ms": { + "median": 0.043904 + }, + "synchronized_e2e_ms": { + "median": 0.0914565 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 39.357289, + "after_init_synchronized_e2e_ms_per_call": 39.386697, + "including_init_host_enqueue_ms_per_call": 76.715312, + "including_init_synchronized_e2e_ms_per_call": 76.82231999999999, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 3.9752425000000002, + "after_init_synchronized_e2e_ms_per_call": 4.02098055, + "including_init_host_enqueue_ms_per_call": 7.711044799999999, + "including_init_synchronized_e2e_ms_per_call": 7.76454285, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.43703785000000006, + "after_init_synchronized_e2e_ms_per_call": 0.48440890499999995, + "including_init_host_enqueue_ms_per_call": 0.81061808, + "including_init_synchronized_e2e_ms_per_call": 0.8587651349999998, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.083217385, + "after_init_synchronized_e2e_ms_per_call": 0.1307517405, + "including_init_host_enqueue_ms_per_call": 0.120575408, + "including_init_synchronized_e2e_ms_per_call": 0.1681873635, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 39.357289, + "synchronized_e2e_ms": 39.386697, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.046688, + "median": 0.046016, + "min": 0.045632, + "p90": 0.046208, + "sample_count": 2174 + }, + "host_enqueue_ms": { + "max": 0.601505, + "median": 0.043904, + "min": 0.035232, + "p90": 0.062092800000000004, + "sample_count": 2174 + }, + "sample_count": 2174, + "synchronized_e2e_ms": { + "max": 0.632993, + "median": 0.0914565, + "min": 0.083296, + "p90": 0.1084064, + "sample_count": 2174 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.066976, + "submission_ms": 0.066976, + "synchronized_e2e_ms": 0.097472 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.044288, + "submission_ms": 0.044288, + "synchronized_e2e_ms": 0.076512 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.193792, + "submission_ms": 0.193792, + "synchronized_e2e_ms": 0.217248 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.225448, + "submission_ms": 7.225448, + "synchronized_e2e_ms": 7.252648 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.06864, + "submission_ms": 0.06864, + "synchronized_e2e_ms": 0.085696 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.042528, + "submission_ms": 0.042528, + "synchronized_e2e_ms": 0.065504 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.223937, + "submission_ms": 1.223937, + "synchronized_e2e_ms": 1.247937 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.058882, + "submission_ms": 1.058882, + "synchronized_e2e_ms": 1.081282 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.069824, + "submission_ms": 0.069824, + "synchronized_e2e_ms": 0.122752 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 39.357289, + "submission_ms": 39.357289, + "synchronized_e2e_ms": 39.386697 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.642463, + "evolution_kernel_ms": 0.193983, + "evolution_speedup": 3.312, + "evolution_tflops": 17.9895, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d416_b2_n2048_k1024_d416", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 9390, + "measurement_schedule_sha256": "37b8bc793a3c8d4bbe4a382e752f0942338ea167f68a2705f01a76dc065175d0", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 2521, + "public_pair_count": 2174, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 2521, + "baseline_public_raw": 2174, + "candidate_precomputed": 2521, + "candidate_public_raw": 2174 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 1880 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:post_d895_d416_b2_n2048_k1024_d416", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.13995840443686, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.8808613931213198, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.18413953320330467, + "including_init_synchronized_e2e_speedup": 6.018614993663301, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.21887201120632133, + "including_init_synchronized_e2e_speedup": 5.974751237286302, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.501277551864989, + "including_init_synchronized_e2e_speedup": 5.58236042034939, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.3697533227100713, + "including_init_synchronized_e2e_speedup": 3.770847802129915, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.623087621696801, + "hot_synchronized_e2e_speedup": 1.8808613931213198, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 241601, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d416_b2_n2048_k1024_d416", + "source": "new_high_gap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 416, + "K": 1024, + "N": 8192, + "baseline_07cf_adapter_bench_iters": 815, + "baseline_07cf_adapter_gpu_span_ms": 0.176352, + "baseline_07cf_adapter_host_enqueue_ms": 0.145824, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.00224, + "baseline_07cf_adapter_kernel_sum_ms": 0.173921, + "baseline_07cf_adapter_submission_ms": 0.145824, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.275457, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.173921 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.188992, + "synchronized_e2e_ms": 0.30368 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.176352 + }, + "host_enqueue_ms": { + "median": 0.145824 + }, + "inter_kernel_gap_ms": { + "median": 0.00224 + }, + "kernel_sum_ms": { + "median": 0.173921 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 815, + "submission_ms": { + "median": 0.145824 + }, + "synchronized_e2e_ms": { + "median": 0.275457 + } + }, + "baseline_07cf_precomputed_bench_iters": 1379, + "baseline_07cf_precomputed_gpu_span_ms": 0.130912, + "baseline_07cf_precomputed_host_enqueue_ms": 0.041632, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.130912, + "baseline_07cf_precomputed_submission_ms": 0.041632, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.1792, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.130912 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.044992, + "synchronized_e2e_ms": 0.178944 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.130912 + }, + "host_enqueue_ms": { + "median": 0.041632 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.130912 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1379, + "submission_ms": { + "median": 0.041632 + }, + "synchronized_e2e_ms": { + "median": 0.1792 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.3471033977022733, + "submission": 3.5026902382782477, + "synchronized_e2e": 1.5371484375 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.223913, + "after_init_synchronized_e2e_ms_per_call": 8.315465, + "including_init_host_enqueue_ms_per_call": 42.45182, + "including_init_synchronized_e2e_ms_per_call": 42.660908000000006, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9536329, + "after_init_synchronized_e2e_ms_per_call": 1.0794578, + "including_init_host_enqueue_ms_per_call": 4.3764236, + "including_init_synchronized_e2e_ms_per_call": 4.514002100000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22660488999999998, + "after_init_synchronized_e2e_ms_per_call": 0.35585708, + "including_init_host_enqueue_ms_per_call": 0.56888396, + "including_init_synchronized_e2e_ms_per_call": 0.69931151, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15390208900000002, + "after_init_synchronized_e2e_ms_per_call": 0.28349700800000005, + "including_init_host_enqueue_ms_per_call": 0.188129996, + "including_init_synchronized_e2e_ms_per_call": 0.317842451, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.223913, + "synchronized_e2e_ms": 8.315465, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.209152, + "median": 0.176352, + "min": 0.168737, + "p90": 0.17886760000000002, + "sample_count": 815 + }, + "host_enqueue_ms": { + "max": 0.391361, + "median": 0.145824, + "min": 0.125472, + "p90": 0.1762304, + "sample_count": 815 + }, + "sample_count": 815, + "synchronized_e2e_ms": { + "max": 0.421441, + "median": 0.275457, + "min": 0.259648, + "p90": 0.2979136, + "sample_count": 815 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 1379, + "candidate_precomputed_gpu_span_ms": 0.072864, + "candidate_precomputed_host_enqueue_ms": 0.051936, + "candidate_precomputed_inter_kernel_gap_ms": 0.002272, + "candidate_precomputed_kernel_sum_ms": 0.07056, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.051936, + "candidate_precomputed_synchronized_e2e_ms": 0.106592, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.07056 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.042048, + "synchronized_e2e_ms": 0.10032 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.072864 + }, + "host_enqueue_ms": { + "median": 0.051936 + }, + "inter_kernel_gap_ms": { + "median": 0.002272 + }, + "kernel_sum_ms": { + "median": 0.07056 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1379, + "submission_ms": { + "median": 0.051936 + }, + "synchronized_e2e_ms": { + "median": 0.106592 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282af3b0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04774d40" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.6864295125164692, + "submission": 0.8441158348736906, + "synchronized_e2e": 1.5869108375863104 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 815, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.12288, + "candidate_public_raw_host_enqueue_ms": 0.04384, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.122912, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.04384, + "candidate_public_raw_synchronized_e2e_ms": 0.169152, + "candidate_public_raw_tflops_from_gpu_span": 227.19146666666666, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.122816 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.049664, + "synchronized_e2e_ms": 0.149984 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.12288 + }, + "host_enqueue_ms": { + "median": 0.04384 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.122912 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 815, + "submission_ms": { + "median": 0.04384 + }, + "synchronized_e2e_ms": { + "median": 0.169152 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.382083, + "after_init_synchronized_e2e_ms_per_call": 2.474531, + "including_init_host_enqueue_ms_per_call": 36.959623, + "including_init_synchronized_e2e_ms_per_call": 450.562323, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.2776643, + "after_init_synchronized_e2e_ms_per_call": 0.3996899, + "including_init_host_enqueue_ms_per_call": 3.7354183, + "including_init_synchronized_e2e_ms_per_call": 45.208469099999995, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.06722243, + "after_init_synchronized_e2e_ms_per_call": 0.19220578999999996, + "including_init_host_enqueue_ms_per_call": 0.41299782999999995, + "including_init_synchronized_e2e_ms_per_call": 4.673083709999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.046178243, + "after_init_synchronized_e2e_ms_per_call": 0.171457379, + "including_init_host_enqueue_ms_per_call": 0.08075578300000001, + "including_init_synchronized_e2e_ms_per_call": 0.619545171, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.382083, + "synchronized_e2e_ms": 2.474531, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.124544, + "median": 0.12288, + "min": 0.12112, + "p90": 0.12384, + "sample_count": 815 + }, + "host_enqueue_ms": { + "max": 0.094464, + "median": 0.04384, + "min": 0.037024, + "p90": 0.0551808, + "sample_count": 815 + }, + "sample_count": 815, + "synchronized_e2e_ms": { + "max": 0.214368, + "median": 0.169152, + "min": 0.16192, + "p90": 0.1783744, + "sample_count": 815 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.063616, + "submission_ms": 0.063616, + "synchronized_e2e_ms": 0.195936 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.044992, + "submission_ms": 0.044992, + "synchronized_e2e_ms": 0.178944 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.188992, + "submission_ms": 0.188992, + "synchronized_e2e_ms": 0.30368 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.223913, + "submission_ms": 8.223913, + "synchronized_e2e_ms": 8.315465 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.065824, + "submission_ms": 0.065824, + "synchronized_e2e_ms": 0.113248 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.042048, + "submission_ms": 0.042048, + "synchronized_e2e_ms": 0.10032 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.248801, + "submission_ms": 1.248801, + "synchronized_e2e_ms": 1.273089 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.07069, + "submission_ms": 1.07069, + "synchronized_e2e_ms": 1.092738 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.049664, + "submission_ms": 0.049664, + "synchronized_e2e_ms": 0.149984 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.382083, + "submission_ms": 2.382083, + "synchronized_e2e_ms": 2.474531 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.710015, + "evolution_kernel_ms": 0.230463, + "evolution_speedup": 3.0808, + "evolution_tflops": 121.1354, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d416_b4_n8192_k1024_d416", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 4388, + "measurement_schedule_sha256": "b85531cb7bd359bddd608e6d902885c6a69b45ff4d6375da06f1ac794e519635", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 1379, + "public_pair_count": 815, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 1379, + "baseline_public_raw": 815, + "candidate_precomputed": 1379, + "candidate_public_raw": 815 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 878 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:post_d895_d416_b4_n8192_k1024_d416", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.7966622749231445, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.6284584279228151, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.360420621119719, + "including_init_synchronized_e2e_speedup": 0.0946837003945401, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.7007382473262393, + "including_init_synchronized_e2e_speedup": 0.09984859451920705, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 1.851437878120113, + "including_init_synchronized_e2e_speedup": 0.14964669015098814, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.6534546932506187, + "including_init_synchronized_e2e_speedup": 0.5130254675167181, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.4351562500000001, + "hot_synchronized_e2e_speedup": 1.6284584279228151, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 241602, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d416_b4_n8192_k1024_d416", + "source": "new_high_gap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 480, + "K": 256, + "N": 256, + "baseline_07cf_adapter_bench_iters": 5522, + "baseline_07cf_adapter_gpu_span_ms": 0.058032, + "baseline_07cf_adapter_host_enqueue_ms": 0.15627200000000002, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.044192, + "baseline_07cf_adapter_kernel_sum_ms": 0.013824, + "baseline_07cf_adapter_submission_ms": 0.15627200000000002, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.177408, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.013824 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.208545, + "synchronized_e2e_ms": 0.229633 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.058032 + }, + "host_enqueue_ms": { + "median": 0.15627200000000002 + }, + "inter_kernel_gap_ms": { + "median": 0.044192 + }, + "kernel_sum_ms": { + "median": 0.013824 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 5522, + "submission_ms": { + "median": 0.15627200000000002 + }, + "synchronized_e2e_ms": { + "median": 0.177408 + } + }, + "baseline_07cf_precomputed_bench_iters": 7567, + "baseline_07cf_precomputed_gpu_span_ms": 0.013184, + "baseline_07cf_precomputed_host_enqueue_ms": 0.045888, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.013184, + "baseline_07cf_precomputed_submission_ms": 0.045888, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.066433, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.013184 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.042048, + "synchronized_e2e_ms": 0.056928 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.013184 + }, + "host_enqueue_ms": { + "median": 0.045888 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.013184 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 7567, + "submission_ms": { + "median": 0.045888 + }, + "synchronized_e2e_ms": { + "median": 0.066433 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 4.401699029126214, + "submission": 3.4055090655509073, + "synchronized_e2e": 2.6704800325139613 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.787305, + "after_init_synchronized_e2e_ms_per_call": 8.814633, + "including_init_host_enqueue_ms_per_call": 43.236749, + "including_init_synchronized_e2e_ms_per_call": 492.856127, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 1.0193753, + "after_init_synchronized_e2e_ms_per_call": 1.0411305, + "including_init_host_enqueue_ms_per_call": 4.4643197, + "including_init_synchronized_e2e_ms_per_call": 49.4452799, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.24258233000000004, + "after_init_synchronized_e2e_ms_per_call": 0.26378025, + "including_init_host_enqueue_ms_per_call": 0.58707677, + "including_init_synchronized_e2e_ms_per_call": 5.10419519, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.16490303300000003, + "after_init_synchronized_e2e_ms_per_call": 0.186045225, + "including_init_host_enqueue_ms_per_call": 0.19935247700000003, + "including_init_synchronized_e2e_ms_per_call": 0.670086719, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.787305, + "synchronized_e2e_ms": 8.814633, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.675777, + "median": 0.058032, + "min": 0.046976, + "p90": 0.084032, + "sample_count": 5522 + }, + "host_enqueue_ms": { + "max": 84.064695, + "median": 0.15627200000000002, + "min": 0.128993, + "p90": 0.2216969, + "sample_count": 5522 + }, + "sample_count": 5522, + "synchronized_e2e_ms": { + "max": 84.295928, + "median": 0.177408, + "min": 0.147201, + "p90": 0.2483849, + "sample_count": 5522 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 7567, + "candidate_precomputed_gpu_span_ms": 0.02192, + "candidate_precomputed_host_enqueue_ms": 0.058656, + "candidate_precomputed_inter_kernel_gap_ms": 0.007168, + "candidate_precomputed_kernel_sum_ms": 0.014784, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.058656, + "candidate_precomputed_synchronized_e2e_ms": 0.071904, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.014784 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.04016, + "synchronized_e2e_ms": 0.05408 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.02192 + }, + "host_enqueue_ms": { + "median": 0.058656 + }, + "inter_kernel_gap_ms": { + "median": 0.007168 + }, + "kernel_sum_ms": { + "median": 0.014784 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 7567, + "submission_ms": { + "median": 0.058656 + }, + "synchronized_e2e_ms": { + "median": 0.071904 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282c6300", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7d3ac530" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 0.8481751824817519, + "submission": 0.8319694489907257, + "synchronized_e2e": 0.9670672007120605 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 5522, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.018592, + "candidate_public_raw_host_enqueue_ms": 0.0488, + "candidate_public_raw_inter_kernel_gap_ms": 3.2e-05, + "candidate_public_raw_kernel_sum_ms": 0.018688, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.0488, + "candidate_public_raw_synchronized_e2e_ms": 0.069536, + "candidate_public_raw_tflops_from_gpu_span": 3.383958691910499, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.01856 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.07712, + "synchronized_e2e_ms": 0.095296 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.018592 + }, + "host_enqueue_ms": { + "median": 0.0488 + }, + "inter_kernel_gap_ms": { + "median": 3.2e-05 + }, + "kernel_sum_ms": { + "median": 0.018688 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 5522, + "submission_ms": { + "median": 0.0488 + }, + "synchronized_e2e_ms": { + "median": 0.069536 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.305795, + "after_init_synchronized_e2e_ms_per_call": 3.333507, + "including_init_host_enqueue_ms_per_call": 38.179879, + "including_init_synchronized_e2e_ms_per_call": 38.276582999999995, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.3744995, + "after_init_synchronized_e2e_ms_per_call": 0.39593310000000004, + "including_init_host_enqueue_ms_per_call": 3.8619079, + "including_init_synchronized_e2e_ms_per_call": 3.8902406999999997, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.08136994999999998, + "after_init_synchronized_e2e_ms_per_call": 0.10217570999999999, + "including_init_host_enqueue_ms_per_call": 0.43011079, + "including_init_synchronized_e2e_ms_per_call": 0.45160647, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.052056995, + "after_init_synchronized_e2e_ms_per_call": 0.072799971, + "including_init_host_enqueue_ms_per_call": 0.08693107900000001, + "including_init_synchronized_e2e_ms_per_call": 0.107743047, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.305795, + "synchronized_e2e_ms": 3.333507, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.019296, + "median": 0.018592, + "min": 0.017888, + "p90": 0.018848, + "sample_count": 5522 + }, + "host_enqueue_ms": { + "max": 25.679706, + "median": 0.0488, + "min": 0.035488, + "p90": 0.07308480000000002, + "sample_count": 5522 + }, + "sample_count": 5522, + "synchronized_e2e_ms": { + "max": 30.858815, + "median": 0.069536, + "min": 0.057984, + "p90": 0.09702720000000033, + "sample_count": 5522 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.060224, + "submission_ms": 0.060224, + "synchronized_e2e_ms": 0.075936 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.042048, + "submission_ms": 0.042048, + "synchronized_e2e_ms": 0.056928 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.208545, + "submission_ms": 0.208545, + "synchronized_e2e_ms": 0.229633 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.787305, + "submission_ms": 8.787305, + "synchronized_e2e_ms": 8.814633 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.08176, + "submission_ms": 0.08176, + "synchronized_e2e_ms": 0.098464 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.04016, + "submission_ms": 0.04016, + "synchronized_e2e_ms": 0.05408 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.328833, + "submission_ms": 1.328833, + "synchronized_e2e_ms": 1.353697 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.073793, + "submission_ms": 1.073793, + "synchronized_e2e_ms": 1.096577 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.07712, + "submission_ms": 0.07712, + "synchronized_e2e_ms": 0.095296 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.305795, + "submission_ms": 3.305795, + "synchronized_e2e_ms": 3.333507 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.649151, + "evolution_kernel_ms": 0.176655, + "evolution_speedup": 3.6747, + "evolution_tflops": 0.3561, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d480_b1_n256_k256_d480", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 26178, + "measurement_schedule_sha256": "fb9c5393bf8ed90c357240dd8f4aaa5b35e3fa9fe5b3fab6b055c4897f63d5b1", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 7567, + "public_pair_count": 5522, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 7567, + "baseline_public_raw": 5522, + "candidate_precomputed": 7567, + "candidate_public_raw": 5522 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 5238 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:post_d895_d480_b1_n256_k256_d480", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.6014598540145986, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.5513115508513575, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.6442521344637946, + "including_init_synchronized_e2e_speedup": 12.8761788114681, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.629561660795725, + "including_init_synchronized_e2e_speedup": 12.710082412124269, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.581633638758175, + "including_init_synchronized_e2e_speedup": 11.302307493513103, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 2.555567295486972, + "including_init_synchronized_e2e_speedup": 6.219303589956947, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 3.121342512908778, + "hot_synchronized_e2e_speedup": 2.5513115508513575, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 248004, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d480_b1_n256_k256_d480", + "source": "high_small_boundary", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 480, + "K": 8192, + "N": 512, + "baseline_07cf_adapter_bench_iters": 2014, + "baseline_07cf_adapter_gpu_span_ms": 0.274368, + "baseline_07cf_adapter_host_enqueue_ms": 0.17521599999999998, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.0385925, + "baseline_07cf_adapter_kernel_sum_ms": 0.235776, + "baseline_07cf_adapter_submission_ms": 0.17521599999999998, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.393024, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.235776 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.239008, + "synchronized_e2e_ms": 0.449632 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.274368 + }, + "host_enqueue_ms": { + "median": 0.17521599999999998 + }, + "inter_kernel_gap_ms": { + "median": 0.0385925 + }, + "kernel_sum_ms": { + "median": 0.235776 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2014, + "submission_ms": { + "median": 0.17521599999999998 + }, + "synchronized_e2e_ms": { + "median": 0.393024 + } + }, + "baseline_07cf_precomputed_bench_iters": 2636, + "baseline_07cf_precomputed_gpu_span_ms": 0.250401, + "baseline_07cf_precomputed_host_enqueue_ms": 0.0542725, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.250401, + "baseline_07cf_precomputed_submission_ms": 0.0542725, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.311424, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.250401 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.06976, + "synchronized_e2e_ms": 0.272448 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.250401 + }, + "host_enqueue_ms": { + "median": 0.0542725 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.250401 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2636, + "submission_ms": { + "median": 0.0542725 + }, + "synchronized_e2e_ms": { + "median": 0.311424 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.0957144739837301, + "submission": 3.228449030356073, + "synchronized_e2e": 1.2620221948212085 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.555049, + "after_init_synchronized_e2e_ms_per_call": 8.748105, + "including_init_host_enqueue_ms_per_call": 44.106317000000004, + "including_init_synchronized_e2e_ms_per_call": 44.379342, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 1.0131993, + "after_init_synchronized_e2e_ms_per_call": 1.2285321, + "including_init_host_enqueue_ms_per_call": 4.5683261, + "including_init_synchronized_e2e_ms_per_call": 4.7916558, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.25901432999999996, + "after_init_synchronized_e2e_ms_per_call": 0.47657481, + "including_init_host_enqueue_ms_per_call": 0.6145270100000001, + "including_init_synchronized_e2e_ms_per_call": 0.8328871800000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.18359583299999996, + "after_init_synchronized_e2e_ms_per_call": 0.40137908099999997, + "including_init_host_enqueue_ms_per_call": 0.21914710099999996, + "including_init_synchronized_e2e_ms_per_call": 0.437010318, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.555049, + "synchronized_e2e_ms": 8.748105, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.607905, + "median": 0.274368, + "min": 0.263457, + "p90": 0.2937703, + "sample_count": 2014 + }, + "host_enqueue_ms": { + "max": 1.038401, + "median": 0.17521599999999998, + "min": 0.144864, + "p90": 0.2123111, + "sample_count": 2014 + }, + "sample_count": 2014, + "synchronized_e2e_ms": { + "max": 1.183457, + "median": 0.393024, + "min": 0.366688, + "p90": 0.42652510000000005, + "sample_count": 2014 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 2636, + "candidate_precomputed_gpu_span_ms": 0.03808, + "candidate_precomputed_host_enqueue_ms": 0.061696, + "candidate_precomputed_inter_kernel_gap_ms": 0.002144, + "candidate_precomputed_kernel_sum_ms": 0.035936, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.061696, + "candidate_precomputed_synchronized_e2e_ms": 0.08203250000000001, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.035936 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.061409, + "synchronized_e2e_ms": 0.078081 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.03808 + }, + "host_enqueue_ms": { + "median": 0.061696 + }, + "inter_kernel_gap_ms": { + "median": 0.002144 + }, + "kernel_sum_ms": { + "median": 0.035936 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2636, + "submission_ms": { + "median": 0.061696 + }, + "synchronized_e2e_ms": { + "median": 0.08203250000000001 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb933d7620", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb933d65a0" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.3042016806722687, + "submission": 1.0534313407676348, + "synchronized_e2e": 1.4138847408039497 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2014, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.049664, + "candidate_public_raw_host_enqueue_ms": 0.0649925, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.0496, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.0649925, + "candidate_public_raw_synchronized_e2e_ms": 0.1159845, + "candidate_public_raw_tflops_from_gpu_span": 81.07546391752577, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.0496 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.279873, + "synchronized_e2e_ms": 0.310305 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.049664 + }, + "host_enqueue_ms": { + "median": 0.0649925 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.0496 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2014, + "submission_ms": { + "median": 0.0649925 + }, + "synchronized_e2e_ms": { + "median": 0.1159845 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.282852, + "after_init_synchronized_e2e_ms_per_call": 3.312196, + "including_init_host_enqueue_ms_per_call": 39.120872, + "including_init_synchronized_e2e_ms_per_call": 487.412409, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.38677845, + "after_init_synchronized_e2e_ms_per_call": 0.43560565, + "including_init_host_enqueue_ms_per_call": 3.97058045, + "including_init_synchronized_e2e_ms_per_call": 48.84562695, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.097171095, + "after_init_synchronized_e2e_ms_per_call": 0.147946615, + "including_init_host_enqueue_ms_per_call": 0.455551295, + "including_init_synchronized_e2e_ms_per_call": 4.988948745, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.0682103595, + "after_init_synchronized_e2e_ms_per_call": 0.1191807115, + "including_init_host_enqueue_ms_per_call": 0.10404837949999998, + "including_init_synchronized_e2e_ms_per_call": 0.6032809245, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.282852, + "synchronized_e2e_ms": 3.312196, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.051264, + "median": 0.049664, + "min": 0.048544, + "p90": 0.0504, + "sample_count": 2014 + }, + "host_enqueue_ms": { + "max": 82.891606, + "median": 0.0649925, + "min": 0.051488, + "p90": 0.08288, + "sample_count": 2014 + }, + "sample_count": 2014, + "synchronized_e2e_ms": { + "max": 83.254583, + "median": 0.1159845, + "min": 0.103488, + "p90": 0.1316384, + "sample_count": 2014 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.082464, + "submission_ms": 0.082464, + "synchronized_e2e_ms": 0.28848 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.06976, + "submission_ms": 0.06976, + "synchronized_e2e_ms": 0.272448 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.239008, + "submission_ms": 0.239008, + "synchronized_e2e_ms": 0.449632 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.555049, + "submission_ms": 8.555049, + "synchronized_e2e_ms": 8.748105 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.100896, + "submission_ms": 0.100896, + "synchronized_e2e_ms": 0.119296 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.061409, + "submission_ms": 0.061409, + "synchronized_e2e_ms": 0.078081 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.390337, + "submission_ms": 1.390337, + "synchronized_e2e_ms": 1.415489 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.180545, + "submission_ms": 1.180545, + "synchronized_e2e_ms": 1.204161 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.279873, + "submission_ms": 0.279873, + "synchronized_e2e_ms": 0.310305 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.282852, + "submission_ms": 3.282852, + "synchronized_e2e_ms": 3.312196 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 32.836689, + "evolution_kernel_ms": 0.278559, + "evolution_speedup": 117.8803, + "evolution_tflops": 14.4548, + "expected_route": "d480_splitk_k1024_eac2_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d480_b1_n512_k8192_d480", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 9300, + "measurement_schedule_sha256": "641beed0ec989afc60b44bbb04bbfb4b7b98e1704a5dfc0e176e81104de4f300", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 2636, + "public_pair_count": 2014, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 2636, + "baseline_public_raw": 2014, + "candidate_precomputed": 2636, + "candidate_public_raw": 2014 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 1862 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:post_d895_d480_b1_n512_k8192_d480", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 6.5756565126050415, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 3.3885907168630287, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.641179749024514, + "including_init_synchronized_e2e_speedup": 0.09105090715899275, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.8202850445121634, + "including_init_synchronized_e2e_speedup": 0.09809794856159584, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 3.2212620072449782, + "including_init_synchronized_e2e_speedup": 0.16694642951277705, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 3.3678191374113418, + "including_init_synchronized_e2e_speedup": 0.7243894183496598, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 5.524484536082475, + "hot_synchronized_e2e_speedup": 3.3885907168630287, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 248003, + "selected_route": "d480_splitk_k1024_eac2_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d480_b1_n512_k8192_d480", + "source": "high_k_low_n", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 480, + "K": 1024, + "N": 2048, + "baseline_07cf_adapter_bench_iters": 2334, + "baseline_07cf_adapter_gpu_span_ms": 0.075776, + "baseline_07cf_adapter_host_enqueue_ms": 0.1436485, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.031888, + "baseline_07cf_adapter_kernel_sum_ms": 0.04384, + "baseline_07cf_adapter_submission_ms": 0.1436485, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.17264, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.04384 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.20512, + "synchronized_e2e_ms": 0.227905 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.075776 + }, + "host_enqueue_ms": { + "median": 0.1436485 + }, + "inter_kernel_gap_ms": { + "median": 0.031888 + }, + "kernel_sum_ms": { + "median": 0.04384 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2334, + "submission_ms": { + "median": 0.1436485 + }, + "synchronized_e2e_ms": { + "median": 0.17264 + } + }, + "baseline_07cf_precomputed_bench_iters": 2688, + "baseline_07cf_precomputed_gpu_span_ms": 0.04384, + "baseline_07cf_precomputed_host_enqueue_ms": 0.04352, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.04384, + "baseline_07cf_precomputed_submission_ms": 0.04352, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.092992, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.04384 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.047968, + "synchronized_e2e_ms": 0.08032 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.04384 + }, + "host_enqueue_ms": { + "median": 0.04352 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.04384 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2688, + "submission_ms": { + "median": 0.04352 + }, + "synchronized_e2e_ms": { + "median": 0.092992 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.7284671532846716, + "submission": 3.3007467830882353, + "synchronized_e2e": 1.8565037852718511 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.509576, + "after_init_synchronized_e2e_ms_per_call": 7.540008, + "including_init_host_enqueue_ms_per_call": 44.444014, + "including_init_synchronized_e2e_ms_per_call": 462.65132700000004, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.8802412500000001, + "after_init_synchronized_e2e_ms_per_call": 0.9093768000000001, + "including_init_host_enqueue_ms_per_call": 4.57368505, + "including_init_synchronized_e2e_ms_per_call": 46.420508700000006, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.217307775, + "after_init_synchronized_e2e_ms_per_call": 0.24631367999999998, + "including_init_host_enqueue_ms_per_call": 0.586652155, + "including_init_synchronized_e2e_ms_per_call": 4.797426870000001, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15101442750000002, + "after_init_synchronized_e2e_ms_per_call": 0.18000736799999997, + "including_init_host_enqueue_ms_per_call": 0.1879488655, + "including_init_synchronized_e2e_ms_per_call": 0.6351186870000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.509576, + "synchronized_e2e_ms": 7.540008, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.434528, + "median": 0.075776, + "min": 0.067808, + "p90": 0.09567040000000002, + "sample_count": 2334 + }, + "host_enqueue_ms": { + "max": 42.131148, + "median": 0.1436485, + "min": 0.124416, + "p90": 0.17615360000000013, + "sample_count": 2334 + }, + "sample_count": 2334, + "synchronized_e2e_ms": { + "max": 42.258284, + "median": 0.17264, + "min": 0.15488, + "p90": 0.2042784, + "sample_count": 2334 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 2688, + "candidate_precomputed_gpu_span_ms": 0.037056, + "candidate_precomputed_host_enqueue_ms": 0.050048, + "candidate_precomputed_inter_kernel_gap_ms": 0.002112, + "candidate_precomputed_kernel_sum_ms": 0.034912, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.050048, + "candidate_precomputed_synchronized_e2e_ms": 0.070816, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.034912 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.047424, + "synchronized_e2e_ms": 0.064416 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.037056 + }, + "host_enqueue_ms": { + "median": 0.050048 + }, + "inter_kernel_gap_ms": { + "median": 0.002112 + }, + "kernel_sum_ms": { + "median": 0.034912 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 2688, + "submission_ms": { + "median": 0.050048 + }, + "synchronized_e2e_ms": { + "median": 0.070816 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc00428c80", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc00428c20" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.151986183074266, + "submission": 0.9341432225063938, + "synchronized_e2e": 1.2828739267962042 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2334, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.042688, + "candidate_public_raw_host_enqueue_ms": 0.046752, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.042688, + "candidate_public_raw_norm_compute_fields": [ + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.046752, + "candidate_public_raw_synchronized_e2e_ms": 0.090848, + "candidate_public_raw_tflops_from_gpu_span": 94.32467766116942, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.042624 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.26096, + "synchronized_e2e_ms": 0.28688 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.042688 + }, + "host_enqueue_ms": { + "median": 0.046752 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.042688 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2334, + "submission_ms": { + "median": 0.046752 + }, + "synchronized_e2e_ms": { + "median": 0.090848 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.481891, + "after_init_synchronized_e2e_ms_per_call": 2.512163, + "including_init_host_enqueue_ms_per_call": 39.83991399999999, + "including_init_synchronized_e2e_ms_per_call": 39.947786, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.29026589999999997, + "after_init_synchronized_e2e_ms_per_call": 0.33297950000000004, + "including_init_host_enqueue_ms_per_call": 4.026068199999999, + "including_init_synchronized_e2e_ms_per_call": 4.0765418, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.07110339, + "after_init_synchronized_e2e_ms_per_call": 0.11506115000000001, + "including_init_host_enqueue_ms_per_call": 0.44468361999999995, + "including_init_synchronized_e2e_ms_per_call": 0.48941738, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.049187139000000005, + "after_init_synchronized_e2e_ms_per_call": 0.093269315, + "including_init_host_enqueue_ms_per_call": 0.08654516200000001, + "including_init_synchronized_e2e_ms_per_call": 0.130704938, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.481891, + "synchronized_e2e_ms": 2.512163, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.044352, + "median": 0.042688, + "min": 0.041568, + "p90": 0.043104, + "sample_count": 2334 + }, + "host_enqueue_ms": { + "max": 0.38208, + "median": 0.046752, + "min": 0.037696, + "p90": 0.06520640000000001, + "sample_count": 2334 + }, + "sample_count": 2334, + "synchronized_e2e_ms": { + "max": 40.623658, + "median": 0.090848, + "min": 0.082336, + "p90": 0.10820480000000002, + "sample_count": 2334 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.068097, + "submission_ms": 0.068097, + "synchronized_e2e_ms": 0.099137 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.047968, + "submission_ms": 0.047968, + "synchronized_e2e_ms": 0.08032 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.20512, + "submission_ms": 0.20512, + "synchronized_e2e_ms": 0.227905 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.509576, + "submission_ms": 7.509576, + "synchronized_e2e_ms": 7.540008 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.07296, + "submission_ms": 0.07296, + "synchronized_e2e_ms": 0.094112 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.047424, + "submission_ms": 0.047424, + "synchronized_e2e_ms": 0.064416 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.341633, + "submission_ms": 1.341633, + "synchronized_e2e_ms": 1.366433 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.282145, + "submission_ms": 1.282145, + "synchronized_e2e_ms": 1.305697 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.26096, + "submission_ms": 0.26096, + "synchronized_e2e_ms": 0.28688 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.481891, + "submission_ms": 2.481891, + "synchronized_e2e_ms": 2.512163 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.672447, + "evolution_kernel_ms": 0.197919, + "evolution_speedup": 3.3976, + "evolution_tflops": 20.3443, + "expected_route": "d480_splitk_k1024_eac2_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d480_b2_n2048_k1024_d480", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 10044, + "measurement_schedule_sha256": "d192f9303d01abbc2ed77403ef3edbbbf4a5d3b2c91dacf9dd00b54bc97fe18c", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 2688, + "public_pair_count": 2334, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 2688, + "baseline_public_raw": 2334, + "candidate_precomputed": 2688, + "candidate_public_raw": 2334 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2010 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:post_d895_d480_b2_n2048_k1024_d480", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.1830742659758204, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.900317013032758, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.001400784901298, + "including_init_synchronized_e2e_speedup": 11.58140095673888, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.7310293876950382, + "including_init_synchronized_e2e_speedup": 11.387227453426334, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.1407197824808804, + "including_init_synchronized_e2e_speedup": 9.802322242826769, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.929974161384159, + "including_init_synchronized_e2e_speedup": 4.85917897761445, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.775112443778111, + "hot_synchronized_e2e_speedup": 1.900317013032758, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 248001, + "selected_route": "d480_splitk_k1024_eac2_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d480_b2_n2048_k1024_d480", + "source": "new_high_gap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 480, + "K": 1024, + "N": 8192, + "baseline_07cf_adapter_bench_iters": 750, + "baseline_07cf_adapter_gpu_span_ms": 0.180896, + "baseline_07cf_adapter_host_enqueue_ms": 0.151712, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.002208, + "baseline_07cf_adapter_kernel_sum_ms": 0.17872, + "baseline_07cf_adapter_submission_ms": 0.151712, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.2827045, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.17872 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.257568, + "synchronized_e2e_ms": 0.36448 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.180896 + }, + "host_enqueue_ms": { + "median": 0.151712 + }, + "inter_kernel_gap_ms": { + "median": 0.002208 + }, + "kernel_sum_ms": { + "median": 0.17872 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 750, + "submission_ms": { + "median": 0.151712 + }, + "synchronized_e2e_ms": { + "median": 0.2827045 + } + }, + "baseline_07cf_precomputed_bench_iters": 1193, + "baseline_07cf_precomputed_gpu_span_ms": 0.13568, + "baseline_07cf_precomputed_host_enqueue_ms": 0.044032, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.13568, + "baseline_07cf_precomputed_submission_ms": 0.044032, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.186497, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.13568 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.05504, + "synchronized_e2e_ms": 0.201505 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.13568 + }, + "host_enqueue_ms": { + "median": 0.044032 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.13568 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1193, + "submission_ms": { + "median": 0.044032 + }, + "synchronized_e2e_ms": { + "median": 0.186497 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.3332547169811322, + "submission": 3.4454941860465116, + "synchronized_e2e": 1.5158662069631148 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.862313, + "after_init_synchronized_e2e_ms_per_call": 8.952169, + "including_init_host_enqueue_ms_per_call": 43.09022, + "including_init_synchronized_e2e_ms_per_call": 43.297612, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 1.0227721, + "after_init_synchronized_e2e_ms_per_call": 1.14965095, + "including_init_host_enqueue_ms_per_call": 4.4455628, + "including_init_synchronized_e2e_ms_per_call": 4.58419525, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.23881801000000002, + "after_init_synchronized_e2e_ms_per_call": 0.369399145, + "including_init_host_enqueue_ms_per_call": 0.58109708, + "including_init_synchronized_e2e_ms_per_call": 0.712853575, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.16042260100000003, + "after_init_synchronized_e2e_ms_per_call": 0.2913739645000001, + "including_init_host_enqueue_ms_per_call": 0.194650508, + "including_init_synchronized_e2e_ms_per_call": 0.32571940750000006, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.862313, + "synchronized_e2e_ms": 8.952169, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.44704, + "median": 0.180896, + "min": 0.172928, + "p90": 0.18336319999999998, + "sample_count": 750 + }, + "host_enqueue_ms": { + "max": 0.439232, + "median": 0.151712, + "min": 0.129536, + "p90": 0.19328640000000002, + "sample_count": 750 + }, + "sample_count": 750, + "synchronized_e2e_ms": { + "max": 0.543392, + "median": 0.2827045, + "min": 0.267008, + "p90": 0.311424, + "sample_count": 750 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.227907, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 34.345443, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 1193, + "candidate_precomputed_gpu_span_ms": 0.08336, + "candidate_precomputed_host_enqueue_ms": 0.053792, + "candidate_precomputed_inter_kernel_gap_ms": 0.002048, + "candidate_precomputed_kernel_sum_ms": 0.081121, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.053792, + "candidate_precomputed_synchronized_e2e_ms": 0.118368, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.081121 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.05024, + "synchronized_e2e_ms": 0.114464 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.08336 + }, + "host_enqueue_ms": { + "median": 0.053792 + }, + "inter_kernel_gap_ms": { + "median": 0.002048 + }, + "kernel_sum_ms": { + "median": 0.081121 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1193, + "submission_ms": { + "median": 0.053792 + }, + "synchronized_e2e_ms": { + "median": 0.118368 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc282ad250", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb9e4d4470" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.6003838771593089, + "submission": 0.8566422516359309, + "synchronized_e2e": 1.5309543119762097 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 750, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.133408, + "candidate_public_raw_host_enqueue_ms": 0.046080499999999996, + "candidate_public_raw_inter_kernel_gap_ms": 9.6e-05, + "candidate_public_raw_kernel_sum_ms": 0.133312, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.046080499999999996, + "candidate_public_raw_synchronized_e2e_ms": 0.181216, + "candidate_public_raw_tflops_from_gpu_span": 241.45669465099544, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.133312 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.068128, + "synchronized_e2e_ms": 0.172737 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.133408 + }, + "host_enqueue_ms": { + "median": 0.046080499999999996 + }, + "inter_kernel_gap_ms": { + "median": 9.6e-05 + }, + "kernel_sum_ms": { + "median": 0.133312 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 750, + "submission_ms": { + "median": 0.046080499999999996 + }, + "synchronized_e2e_ms": { + "median": 0.181216 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 3.014531, + "after_init_synchronized_e2e_ms_per_call": 3.109571, + "including_init_host_enqueue_ms_per_call": 37.592071, + "including_init_synchronized_e2e_ms_per_call": 451.197363, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.34292555, + "after_init_synchronized_e2e_ms_per_call": 0.4740515, + "including_init_host_enqueue_ms_per_call": 3.8006795499999995, + "including_init_synchronized_e2e_ms_per_call": 45.2828307, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.075765005, + "after_init_synchronized_e2e_ms_per_call": 0.21049954999999998, + "including_init_host_enqueue_ms_per_call": 0.42154040499999995, + "including_init_synchronized_e2e_ms_per_call": 4.69137747, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.049048950499999994, + "after_init_synchronized_e2e_ms_per_call": 0.18414435499999998, + "including_init_host_enqueue_ms_per_call": 0.08362649049999998, + "including_init_synchronized_e2e_ms_per_call": 0.632232147, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 3.014531, + "synchronized_e2e_ms": 3.109571, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.134881, + "median": 0.133408, + "min": 0.131904, + "p90": 0.134112, + "sample_count": 750 + }, + "host_enqueue_ms": { + "max": 0.106944, + "median": 0.046080499999999996, + "min": 0.037376, + "p90": 0.0629216, + "sample_count": 750 + }, + "sample_count": 750, + "synchronized_e2e_ms": { + "max": 0.228705, + "median": 0.181216, + "min": 0.173472, + "p90": 0.1944128, + "sample_count": 750 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.57754, + "sample_id": "a591a388bda74a2890fb7b17b61b2608", + "synchronized_e2e_ms": 448.087792, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.085696, + "submission_ms": 0.085696, + "synchronized_e2e_ms": 0.222688 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.05504, + "submission_ms": 0.05504, + "synchronized_e2e_ms": 0.201505 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.257568, + "submission_ms": 0.257568, + "synchronized_e2e_ms": 0.36448 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.862313, + "submission_ms": 8.862313, + "synchronized_e2e_ms": 8.952169 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.08496, + "submission_ms": 0.08496, + "synchronized_e2e_ms": 0.137056 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.05024, + "submission_ms": 0.05024, + "synchronized_e2e_ms": 0.114464 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.571617, + "submission_ms": 1.571617, + "synchronized_e2e_ms": 1.605569 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.420706, + "submission_ms": 1.420706, + "synchronized_e2e_ms": 1.450274 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.068128, + "submission_ms": 0.068128, + "synchronized_e2e_ms": 0.172737 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 3.014531, + "submission_ms": 3.014531, + "synchronized_e2e_ms": 3.109571 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.733983, + "evolution_kernel_ms": 0.24, + "evolution_speedup": 3.0583, + "evolution_tflops": 134.2177, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d480_b4_n8192_k1024_d480", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "a591a388bda74a2890fb7b17b61b2608", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 3886, + "measurement_schedule_sha256": "bb05c727b4ea5a9d237839cf9025aa9fb9e5bfc498d001c3381a42cb4f90c646", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 1193, + "public_pair_count": 750, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 1193, + "baseline_public_raw": 750, + "candidate_precomputed": 1193, + "candidate_public_raw": 750 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 778 + }, + "paired_cupti_session_id": "a591a388bda74a2890fb7b17b61b2608:post_d895_d480_b4_n8192_k1024_d480", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.6276391554702494, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.5600416078050505, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 2.878908055162593, + "including_init_synchronized_e2e_speedup": 0.09596158034283547, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.4251604519762093, + "including_init_synchronized_e2e_speedup": 0.10123473243910963, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 1.7548690484136429, + "including_init_synchronized_e2e_speedup": 0.15194973748296575, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.5823127703263025, + "including_init_synchronized_e2e_speedup": 0.5151895692833223, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.3559606620292637, + "hot_synchronized_e2e_speedup": 1.5600416078050505, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 248002, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d480_b4_n8192_k1024_d480", + "source": "new_high_gap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 512, + "kernel": "triton_h200_small_d", + "num_stages": 2, + "num_warps": 4 + } + }, + { + "B": 1, + "D": 48, + "K": 8192, + "N": 512, + "baseline_07cf_adapter_bench_iters": 1714, + "baseline_07cf_adapter_gpu_span_ms": 0.1256, + "baseline_07cf_adapter_host_enqueue_ms": 0.149888, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.039168, + "baseline_07cf_adapter_kernel_sum_ms": 0.086464, + "baseline_07cf_adapter_submission_ms": 0.149888, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.228049, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.086464 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.231424, + "synchronized_e2e_ms": 0.301856 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.1256 + }, + "host_enqueue_ms": { + "median": 0.149888 + }, + "inter_kernel_gap_ms": { + "median": 0.039168 + }, + "kernel_sum_ms": { + "median": 0.086464 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1714, + "submission_ms": { + "median": 0.149888 + }, + "synchronized_e2e_ms": { + "median": 0.228049 + } + }, + "baseline_07cf_precomputed_bench_iters": 1478, + "baseline_07cf_precomputed_gpu_span_ms": 0.091457, + "baseline_07cf_precomputed_host_enqueue_ms": 0.045568, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.091457, + "baseline_07cf_precomputed_submission_ms": 0.045568, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.144272, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.091457 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.054016, + "synchronized_e2e_ms": 0.13456 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.091457 + }, + "host_enqueue_ms": { + "median": 0.045568 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.091457 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1478, + "submission_ms": { + "median": 0.045568 + }, + "synchronized_e2e_ms": { + "median": 0.144272 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 1.3733229823851645, + "submission": 3.289325842696629, + "synchronized_e2e": 1.58068786736165 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.778856, + "after_init_synchronized_e2e_ms_per_call": 7.829609, + "including_init_host_enqueue_ms_per_call": 42.2283, + "including_init_synchronized_e2e_ms_per_call": 491.871103, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9127848000000001, + "after_init_synchronized_e2e_ms_per_call": 0.988205, + "including_init_host_enqueue_ms_per_call": 4.3577292, + "including_init_synchronized_e2e_ms_per_call": 49.3923544, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.22617768, + "after_init_synchronized_e2e_ms_per_call": 0.3040646, + "including_init_host_enqueue_ms_per_call": 0.57067212, + "including_init_synchronized_e2e_ms_per_call": 5.14447954, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15751696799999998, + "after_init_synchronized_e2e_ms_per_call": 0.23565056, + "including_init_host_enqueue_ms_per_call": 0.191966412, + "including_init_synchronized_e2e_ms_per_call": 0.719692054, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.778856, + "synchronized_e2e_ms": 7.829609, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.417504, + "median": 0.1256, + "min": 0.118368, + "p90": 0.1313728, + "sample_count": 1714 + }, + "host_enqueue_ms": { + "max": 0.483809, + "median": 0.149888, + "min": 0.130432, + "p90": 0.1723104, + "sample_count": 1714 + }, + "sample_count": 1714, + "synchronized_e2e_ms": { + "max": 0.549601, + "median": 0.228049, + "min": 0.21024, + "p90": 0.2494848, + "sample_count": 1714 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.449444, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 484.041494, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 1478, + "candidate_precomputed_gpu_span_ms": 0.064512, + "candidate_precomputed_host_enqueue_ms": 0.057824, + "candidate_precomputed_inter_kernel_gap_ms": 0.006816, + "candidate_precomputed_kernel_sum_ms": 0.057632, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.057824, + "candidate_precomputed_synchronized_e2e_ms": 0.1024325, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.057632 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.049312, + "synchronized_e2e_ms": 0.09152 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.064512 + }, + "host_enqueue_ms": { + "median": 0.057824 + }, + "inter_kernel_gap_ms": { + "median": 0.006816 + }, + "kernel_sum_ms": { + "median": 0.057632 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 1478, + "submission_ms": { + "median": 0.057824 + }, + "synchronized_e2e_ms": { + "median": 0.1024325 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7316ee40", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb7316d970" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 0.9077380952380952, + "submission": 0.8074156059767571, + "synchronized_e2e": 1.0527957435384279 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 1714, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.05856, + "candidate_public_raw_host_enqueue_ms": 0.046688, + "candidate_public_raw_inter_kernel_gap_ms": 0.0, + "candidate_public_raw_kernel_sum_ms": 0.058784, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.046688, + "candidate_public_raw_synchronized_e2e_ms": 0.1078405, + "candidate_public_raw_tflops_from_gpu_span": 6.875908196721311, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.05856 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.108608, + "synchronized_e2e_ms": 0.14816 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.05856 + }, + "host_enqueue_ms": { + "median": 0.046688 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.058784 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 1714, + "submission_ms": { + "median": 0.046688 + }, + "synchronized_e2e_ms": { + "median": 0.1078405 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 42.080779, + "after_init_synchronized_e2e_ms_per_call": 42.126475, + "including_init_host_enqueue_ms_per_call": 76.95486299999999, + "including_init_synchronized_e2e_ms_per_call": 77.06955099999999, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 4.2500971, + "after_init_synchronized_e2e_ms_per_call": 4.30970395, + "including_init_host_enqueue_ms_per_call": 7.737505499999999, + "including_init_synchronized_e2e_ms_per_call": 7.8040115499999985, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.46702891, + "after_init_synchronized_e2e_ms_per_call": 0.528026845, + "including_init_host_enqueue_ms_per_call": 0.8157697499999998, + "including_init_synchronized_e2e_ms_per_call": 0.8774576049999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.088722091, + "after_init_synchronized_e2e_ms_per_call": 0.14985913450000002, + "including_init_host_enqueue_ms_per_call": 0.12359617499999999, + "including_init_synchronized_e2e_ms_per_call": 0.1848022105, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 42.080779, + "synchronized_e2e_ms": 42.126475, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.059137, + "median": 0.05856, + "min": 0.058144, + "p90": 0.058816, + "sample_count": 1714 + }, + "host_enqueue_ms": { + "max": 0.0968, + "median": 0.046688, + "min": 0.035424, + "p90": 0.05419840000000001, + "sample_count": 1714 + }, + "sample_count": 1714, + "synchronized_e2e_ms": { + "max": 0.167072, + "median": 0.1078405, + "min": 0.097984, + "p90": 0.11524510000000002, + "sample_count": 1714 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 34.874083999999996, + "sample_id": "41892d8f27794f7aba63e10273932712", + "synchronized_e2e_ms": 34.943076, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.106561, + "submission_ms": 0.106561, + "synchronized_e2e_ms": 0.186177 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.054016, + "submission_ms": 0.054016, + "synchronized_e2e_ms": 0.13456 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.231424, + "submission_ms": 0.231424, + "synchronized_e2e_ms": 0.301856 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.778856, + "submission_ms": 7.778856, + "synchronized_e2e_ms": 7.829609 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.091872, + "submission_ms": 0.091872, + "synchronized_e2e_ms": 0.12112 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.049312, + "submission_ms": 0.049312, + "synchronized_e2e_ms": 0.09152 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.741058, + "submission_ms": 1.741058, + "synchronized_e2e_ms": 1.770562 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.315905, + "submission_ms": 1.315905, + "synchronized_e2e_ms": 1.341505 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.108608, + "submission_ms": 0.108608, + "synchronized_e2e_ms": 0.14816 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 42.080779, + "submission_ms": 42.080779, + "synchronized_e2e_ms": 42.126475 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 1.464301, + "evolution_kernel_ms": 0.21952, + "evolution_speedup": 6.6705, + "evolution_tflops": 1.8342, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d48_b1_n512_k8192_d48", + "measurement_order": [ + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "41892d8f27794f7aba63e10273932712", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "candidate_precomputed", + "baseline_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 6384, + "measurement_schedule_sha256": "d17a599e0e27853c281f5371e8d93dbd0b83ef04174f7132340b158b1435cf31", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 1478, + "public_pair_count": 1714, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 1478, + "baseline_public_raw": 1714, + "candidate_precomputed": 1478, + "candidate_public_raw": 1714 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 1278 + }, + "paired_cupti_session_id": "41892d8f27794f7aba63e10273932712:post_d895_d48_b1_n512_k8192_d48", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 1.4176742311507935, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.1146878955494457, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.18585958117786974, + "including_init_synchronized_e2e_speedup": 6.3821716438960445, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.22929765280048992, + "including_init_synchronized_e2e_speedup": 6.329098064956095, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.5758506463814355, + "including_init_synchronized_e2e_speedup": 5.8629380048509585, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.5724804549701972, + "including_init_synchronized_e2e_speedup": 3.8943909385759214, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 2.1448087431693987, + "hot_synchronized_e2e_speedup": 2.1146878955494457, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 24803, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d48_b1_n512_k8192_d48", + "source": "high_k_low_n", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 128, + "BLOCK_N": 64, + "D_PAD": 64, + "kernel": "triton_h200_small_d", + "num_stages": 4, + "num_warps": 4 + } + }, + { + "B": 2, + "D": 48, + "K": 512, + "N": 2048, + "baseline_07cf_adapter_bench_iters": 7495, + "baseline_07cf_adapter_gpu_span_ms": 0.063552, + "baseline_07cf_adapter_host_enqueue_ms": 0.175041, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.04976, + "baseline_07cf_adapter_kernel_sum_ms": 0.013792, + "baseline_07cf_adapter_submission_ms": 0.175041, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.198433, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.013792 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.266529, + "synchronized_e2e_ms": 0.290561 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.063552 + }, + "host_enqueue_ms": { + "median": 0.175041 + }, + "inter_kernel_gap_ms": { + "median": 0.04976 + }, + "kernel_sum_ms": { + "median": 0.013792 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 7495, + "submission_ms": { + "median": 0.175041 + }, + "synchronized_e2e_ms": { + "median": 0.198433 + } + }, + "baseline_07cf_precomputed_bench_iters": 9273, + "baseline_07cf_precomputed_gpu_span_ms": 0.010912, + "baseline_07cf_precomputed_host_enqueue_ms": 0.054528, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.010912, + "baseline_07cf_precomputed_submission_ms": 0.054528, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.071904, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.010912 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.0584, + "synchronized_e2e_ms": 0.07584 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.010912 + }, + "host_enqueue_ms": { + "median": 0.054528 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.010912 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 9273, + "submission_ms": { + "median": 0.054528 + }, + "synchronized_e2e_ms": { + "median": 0.071904 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 5.8240469208211145, + "submission": 3.210112235915493, + "synchronized_e2e": 2.759693480195817 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 8.424104, + "after_init_synchronized_e2e_ms_per_call": 8.453928, + "including_init_host_enqueue_ms_per_call": 43.975372, + "including_init_synchronized_e2e_ms_per_call": 44.085164999999996, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9999473, + "after_init_synchronized_e2e_ms_per_call": 1.0239825, + "including_init_host_enqueue_ms_per_call": 4.555074100000001, + "including_init_synchronized_e2e_ms_per_call": 4.587106199999999, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.25753163, + "after_init_synchronized_e2e_ms_per_call": 0.28098795000000004, + "including_init_host_enqueue_ms_per_call": 0.61304431, + "including_init_synchronized_e2e_ms_per_call": 0.6373003199999999, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.183290063, + "after_init_synchronized_e2e_ms_per_call": 0.206688495, + "including_init_host_enqueue_ms_per_call": 0.218841331, + "including_init_synchronized_e2e_ms_per_call": 0.24231973199999998, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 8.424104, + "synchronized_e2e_ms": 8.453928, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.526369, + "median": 0.063552, + "min": 0.048896, + "p90": 0.089056, + "sample_count": 7495 + }, + "host_enqueue_ms": { + "max": 1.183425, + "median": 0.175041, + "min": 0.134688, + "p90": 0.2268992, + "sample_count": 7495 + }, + "sample_count": 7495, + "synchronized_e2e_ms": { + "max": 6.01687, + "median": 0.198433, + "min": 0.153568, + "p90": 0.25541800000000003, + "sample_count": 7495 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.551268, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 35.631237, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 9273, + "candidate_precomputed_gpu_span_ms": 0.017248, + "candidate_precomputed_host_enqueue_ms": 0.066912, + "candidate_precomputed_inter_kernel_gap_ms": 0.008064, + "candidate_precomputed_kernel_sum_ms": 0.009152, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.066912, + "candidate_precomputed_synchronized_e2e_ms": 0.080928, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.009152 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.06176, + "synchronized_e2e_ms": 0.078432 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.017248 + }, + "host_enqueue_ms": { + "median": 0.066912 + }, + "inter_kernel_gap_ms": { + "median": 0.008064 + }, + "kernel_sum_ms": { + "median": 0.009152 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 9273, + "submission_ms": { + "median": 0.066912 + }, + "synchronized_e2e_ms": { + "median": 0.080928 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04666e40", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc04667980" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 0.7847866419294991, + "submission": 0.9177427068388331, + "synchronized_e2e": 0.9616449189402927 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 7495, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": true, + "nvrtc_compile_occurred": true, + "scratch_allocation_occurred": true, + "source_read_occurred": true + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.013536, + "candidate_public_raw_host_enqueue_ms": 0.061408, + "candidate_public_raw_inter_kernel_gap_ms": 6.4e-05, + "candidate_public_raw_kernel_sum_ms": 0.013472, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.061408, + "candidate_public_raw_synchronized_e2e_ms": 0.077824, + "candidate_public_raw_tflops_from_gpu_span": 14.873418439716312, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.013472 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.069664, + "synchronized_e2e_ms": 0.091328 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.013536 + }, + "host_enqueue_ms": { + "median": 0.061408 + }, + "inter_kernel_gap_ms": { + "median": 6.4e-05 + }, + "kernel_sum_ms": { + "median": 0.013472 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 7495, + "submission_ms": { + "median": 0.061408 + }, + "synchronized_e2e_ms": { + "median": 0.077824 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 43.603565, + "after_init_synchronized_e2e_ms_per_call": 43.658509, + "including_init_host_enqueue_ms_per_call": 79.441585, + "including_init_synchronized_e2e_ms_per_call": 527.758722, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 4.4156237, + "after_init_synchronized_e2e_ms_per_call": 4.4358925, + "including_init_host_enqueue_ms_per_call": 7.999425700000001, + "including_init_synchronized_e2e_ms_per_call": 52.845913800000005, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.49682957, + "after_init_synchronized_e2e_ms_per_call": 0.5136308500000001, + "including_init_host_enqueue_ms_per_call": 0.85520977, + "including_init_synchronized_e2e_ms_per_call": 5.35463298, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.10495015699999999, + "after_init_synchronized_e2e_ms_per_call": 0.121404685, + "including_init_host_enqueue_ms_per_call": 0.140788177, + "including_init_synchronized_e2e_ms_per_call": 0.6055048980000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 43.603565, + "synchronized_e2e_ms": 43.658509, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.014112, + "median": 0.013536, + "min": 0.013184, + "p90": 0.01376, + "sample_count": 7495 + }, + "host_enqueue_ms": { + "max": 100.039079, + "median": 0.061408, + "min": 0.04224, + "p90": 0.0824, + "sample_count": 7495 + }, + "sample_count": 7495, + "synchronized_e2e_ms": { + "max": 100.382088, + "median": 0.077824, + "min": 0.058656, + "p90": 0.10475520000000002, + "sample_count": 7495 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 35.83802, + "sample_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "synchronized_e2e_ms": 484.100213, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.08064, + "submission_ms": 0.08064, + "synchronized_e2e_ms": 0.101344 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.0584, + "submission_ms": 0.0584, + "synchronized_e2e_ms": 0.07584 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.266529, + "submission_ms": 0.266529, + "synchronized_e2e_ms": 0.290561 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 8.424104, + "submission_ms": 8.424104, + "synchronized_e2e_ms": 8.453928 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.129121, + "submission_ms": 0.129121, + "synchronized_e2e_ms": 0.151265 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.06176, + "submission_ms": 0.06176, + "synchronized_e2e_ms": 0.078432 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.353921, + "submission_ms": 1.353921, + "synchronized_e2e_ms": 1.381217 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.116961, + "submission_ms": 1.116961, + "synchronized_e2e_ms": 1.139777 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.069664, + "submission_ms": 0.069664, + "synchronized_e2e_ms": 0.091328 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 43.603565, + "submission_ms": 43.603565, + "synchronized_e2e_ms": 43.658509 + }, + "cold_measurement_order": [ + "baseline", + "candidate" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.30624, + "evolution_kernel_ms": 0.166656, + "evolution_speedup": 1.8376, + "evolution_tflops": 1.208, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d48_b2_n2048_k512_d48", + "measurement_order": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 33536, + "measurement_schedule_sha256": "cbb675c50d20d4dc6ef3667b7d795843f2f2b019dd0248d25ff4c82a69144d45", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 9273, + "public_pair_count": 7495, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 9273, + "baseline_public_raw": 7495, + "candidate_precomputed": 9273, + "candidate_public_raw": 7495 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 6708 + }, + "paired_cupti_session_id": "958ecd22358d4d5cb8ab46b2ff2cd831:post_d895_d48_b2_n2048_k512_d48", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.6326530612244898, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 2.549766138980263, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 0.19363757933190065, + "including_init_synchronized_e2e_speedup": 0.08353280232477142, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 0.23084024240894027, + "including_init_synchronized_e2e_speedup": 0.0868015305281749, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 0.5470620582856345, + "including_init_synchronized_e2e_speedup": 0.11901848779932625, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.7024754440077827, + "including_init_synchronized_e2e_speedup": 0.40019450346378527, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 4.695035460992908, + "hot_synchronized_e2e_speedup": 2.549766138980263, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 24801, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d48_b2_n2048_k512_d48", + "source": "new_lowmid_gap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 64, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + }, + { + "B": 4, + "D": 48, + "K": 1024, + "N": 8192, + "baseline_07cf_adapter_bench_iters": 2406, + "baseline_07cf_adapter_gpu_span_ms": 0.06856, + "baseline_07cf_adapter_host_enqueue_ms": 0.1431845, + "baseline_07cf_adapter_inter_kernel_gap_ms": 0.02272, + "baseline_07cf_adapter_kernel_sum_ms": 0.045792, + "baseline_07cf_adapter_submission_ms": 0.1431845, + "baseline_07cf_adapter_synchronized_e2e_ms": 0.1660165, + "baseline_07cf_adapter_timing_backend": "cupti", + "baseline_07cf_adapter_timing_diagnostics": { + "active_union_ms": { + "median": 0.045792 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.198304, + "synchronized_e2e_ms": 0.21936 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.06856 + }, + "host_enqueue_ms": { + "median": 0.1431845 + }, + "inter_kernel_gap_ms": { + "median": 0.02272 + }, + "kernel_sum_ms": { + "median": 0.045792 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2406, + "submission_ms": { + "median": 0.1431845 + }, + "synchronized_e2e_ms": { + "median": 0.1660165 + } + }, + "baseline_07cf_precomputed_bench_iters": 4083, + "baseline_07cf_precomputed_gpu_span_ms": 0.024544, + "baseline_07cf_precomputed_host_enqueue_ms": 0.041888, + "baseline_07cf_precomputed_inter_kernel_gap_ms": 0.0, + "baseline_07cf_precomputed_kernel_sum_ms": 0.024544, + "baseline_07cf_precomputed_submission_ms": 0.041888, + "baseline_07cf_precomputed_synchronized_e2e_ms": 0.07232, + "baseline_07cf_precomputed_timing_backend": "cupti", + "baseline_07cf_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.024544 + }, + "activity_count": { + "median": 1.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.045568, + "synchronized_e2e_ms": 0.074624 + }, + "correlated_kernel_activity_count": { + "median": 1.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.024544 + }, + "host_enqueue_ms": { + "median": 0.041888 + }, + "inter_kernel_gap_ms": { + "median": 0.0 + }, + "kernel_sum_ms": { + "median": 0.024544 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4083, + "submission_ms": { + "median": 0.041888 + }, + "synchronized_e2e_ms": { + "median": 0.07232 + } + }, + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "baseline_public_prepared_ratios": { + "gpu_span": 2.7933507170795306, + "submission": 3.418270148968678, + "synchronized_e2e": 2.2955821349557524 + }, + "baseline_public_raw_assignment_launch_count": 1, + "baseline_public_raw_fresh_pointer_cache_hit": true, + "baseline_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "baseline_public_raw_norm_launch_count": 1, + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_shape_first_was_cache_hit": false, + "baseline_public_raw_total_launch_count": 2, + "baseline_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 7.880169, + "after_init_synchronized_e2e_ms_per_call": 7.908201, + "including_init_host_enqueue_ms_per_call": 44.814607, + "including_init_synchronized_e2e_ms_per_call": 463.01952000000006, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.9168829500000001, + "after_init_synchronized_e2e_ms_per_call": 0.94023495, + "including_init_host_enqueue_ms_per_call": 4.6103267500000005, + "including_init_synchronized_e2e_ms_per_call": 46.45136685000001, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.220554345, + "after_init_synchronized_e2e_ms_per_call": 0.243438345, + "including_init_host_enqueue_ms_per_call": 0.5898987250000001, + "including_init_synchronized_e2e_ms_per_call": 4.794551535, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.15092148449999998, + "after_init_synchronized_e2e_ms_per_call": 0.1737586845, + "including_init_host_enqueue_ms_per_call": 0.1878559225, + "including_init_synchronized_e2e_ms_per_call": 0.6288700035000001, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "triton_h200_07cf_raw_adapter_v1.compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 7.880169, + "synchronized_e2e_ms": 7.908201, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.38256, + "median": 0.06856, + "min": 0.060384, + "p90": 0.09051200000000001, + "sample_count": 2406 + }, + "host_enqueue_ms": { + "max": 0.890721, + "median": 0.1431845, + "min": 0.124353, + "p90": 0.1793285, + "sample_count": 2406 + }, + "sample_count": 2406, + "synchronized_e2e_ms": { + "max": 0.953601, + "median": 0.1660165, + "min": 0.146816, + "p90": 0.2057925, + "sample_count": 2406 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 36.934438, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 455.11131900000004, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" + }, + "candidate_graph_capture_error": null, + "candidate_graph_kernel_count": 3, + "candidate_hot_launch_path": "cuda_graph", + "candidate_precomputed_bench_iters": 4083, + "candidate_precomputed_gpu_span_ms": 0.02464, + "candidate_precomputed_host_enqueue_ms": 0.05232, + "candidate_precomputed_inter_kernel_gap_ms": 0.003552, + "candidate_precomputed_kernel_sum_ms": 0.021024, + "candidate_precomputed_launch_count": 2, + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "candidate_precomputed_submission_ms": 0.05232, + "candidate_precomputed_synchronized_e2e_ms": 0.065056, + "candidate_precomputed_timing_backend": "cupti", + "candidate_precomputed_timing_diagnostics": { + "active_union_ms": { + "median": 0.021024 + }, + "activity_count": { + "median": 2.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.04496, + "synchronized_e2e_ms": 0.060896 + }, + "correlated_kernel_activity_count": { + "median": 2.0 + }, + "correlated_launch_activity_count": { + "median": 2.0 + }, + "gpu_span_ms": { + "median": 0.02464 + }, + "host_enqueue_ms": { + "median": 0.05232 + }, + "inter_kernel_gap_ms": { + "median": 0.003552 + }, + "kernel_sum_ms": { + "median": 0.021024 + }, + "lane_primary_metric": "gpu_span_ms", + "sample_count": 4083, + "submission_ms": { + "median": 0.05232 + }, + "synchronized_e2e_ms": { + "median": 0.065056 + } + }, + "candidate_prepared_path_proof": { + "caller_owned_outputs": true, + "internal_device_wide_synchronization_count": 0, + "parent_dispatch_traversal_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + "resolved_direct_launcher": true, + "root_dispatch_traversal_count": 0, + "scratch_reused": true, + "stable_direct_launcher_tokens": [ + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffc028acef0", + "flashlib_cake_kmeans._dispatch_runtime.PreparedKernelSequence@fffb67417710" + ], + "stream_bound": true + }, + "candidate_public_prepared_ratios": { + "gpu_span": 1.6896103896103898, + "submission": 0.8440366972477064, + "synchronized_e2e": 1.337432365961633 + }, + "candidate_public_raw_assignment_launch_count": 2, + "candidate_public_raw_bench_iters": 2406, + "candidate_public_raw_cold_first_call_activity": { + "module_load_occurred": false, + "nvrtc_compile_occurred": false, + "scratch_allocation_occurred": true, + "source_read_occurred": false + }, + "candidate_public_raw_fresh_pointer_cache_hit": true, + "candidate_public_raw_gpu_span_ms": 0.041632, + "candidate_public_raw_host_enqueue_ms": 0.04416, + "candidate_public_raw_inter_kernel_gap_ms": 0.000128, + "candidate_public_raw_kernel_sum_ms": 0.041536, + "candidate_public_raw_norm_compute_fields": [ + "x_sq", + "c_sq" + ], + "candidate_public_raw_norm_launch_count": 1, + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_public_raw_shape_first_was_cache_hit": false, + "candidate_public_raw_submission_ms": 0.04416, + "candidate_public_raw_synchronized_e2e_ms": 0.087008, + "candidate_public_raw_tflops_from_gpu_span": 77.37378631821674, + "candidate_public_raw_timing_backend": "cupti", + "candidate_public_raw_timing_diagnostics": { + "active_union_ms": { + "median": 0.041536 + }, + "activity_count": { + "median": 3.0 + }, + "cold_first_call": { + "host_enqueue_ms": 0.074976, + "synchronized_e2e_ms": 0.103744 + }, + "correlated_kernel_activity_count": { + "median": 3.0 + }, + "correlated_launch_activity_count": { + "median": 1.0 + }, + "gpu_span_ms": { + "median": 0.041632 + }, + "host_enqueue_ms": { + "median": 0.04416 + }, + "inter_kernel_gap_ms": { + "median": 0.000128 + }, + "kernel_sum_ms": { + "median": 0.041536 + }, + "lane_primary_metric": "synchronized_e2e_ms", + "sample_count": 2406, + "submission_ms": { + "median": 0.04416 + }, + "synchronized_e2e_ms": { + "median": 0.087008 + } + }, + "candidate_public_raw_total_launch_count": 3, + "candidate_runtime_lifecycle": { + "amortization": { + "call_counts": [ + { + "after_init_host_enqueue_ms_per_call": 2.39517, + "after_init_synchronized_e2e_ms_per_call": 2.420834, + "including_init_host_enqueue_ms_per_call": 39.753192999999996, + "including_init_synchronized_e2e_ms_per_call": 39.856457, + "public_call_count": 1 + }, + { + "after_init_host_enqueue_ms_per_call": 0.279261, + "after_init_synchronized_e2e_ms_per_call": 0.32039059999999997, + "including_init_host_enqueue_ms_per_call": 4.0150633, + "including_init_synchronized_e2e_ms_per_call": 4.063952899999999, + "public_call_count": 10 + }, + { + "after_init_host_enqueue_ms_per_call": 0.0676701, + "after_init_synchronized_e2e_ms_per_call": 0.11034625999999999, + "including_init_host_enqueue_ms_per_call": 0.44125032999999997, + "including_init_synchronized_e2e_ms_per_call": 0.48470248999999993, + "public_call_count": 100 + }, + { + "after_init_host_enqueue_ms_per_call": 0.04651101, + "after_init_synchronized_e2e_ms_per_call": 0.089341826, + "including_init_host_enqueue_ms_per_call": 0.083869033, + "including_init_synchronized_e2e_ms_per_call": 0.126777449, + "public_call_count": 1000 + } + ], + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "model": "observed_first_call_plus_repeated_hot_median" + }, + "api": "flashlib_cake_kmeans.init(...).compute", + "first_compute": { + "cache_state": "shape_slot_miss", + "code_cache_state": "process_order_dependent", + "gpu_span_ms": null, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + "host_enqueue_ms": 2.39517, + "synchronized_e2e_ms": 2.420834, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "hot_compute": { + "cache_state": "fresh_pointer_shape_slot_hit", + "gpu_span_ms": { + "max": 0.042304, + "median": 0.041632, + "min": 0.041472, + "p90": 0.04192, + "sample_count": 2406 + }, + "host_enqueue_ms": { + "max": 0.530625, + "median": 0.04416, + "min": 0.036448, + "p90": 0.06425600000000001, + "sample_count": 2406 + }, + "sample_count": 2406, + "synchronized_e2e_ms": { + "max": 0.603457, + "median": 0.087008, + "min": 0.08, + "p90": 0.10531199999999999, + "sample_count": 2406 + }, + "timing_backend": "cupti" + }, + "init_once": { + "host_enqueue_ms": 37.358022999999996, + "sample_id": "429dbdfd2f734f02819474b9385e8bb8", + "synchronized_e2e_ms": 37.435623, + "timing_class": "cupti_timestamp_host_diagnostic" + }, + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "schema": "loom-public-runtime-lifecycle-v1", + "timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" + }, + "child_route": null, + "cold_baseline_precomputed_first_pointer": { + "host_enqueue_ms": 0.064896, + "submission_ms": 0.064896, + "synchronized_e2e_ms": 0.093184 + }, + "cold_baseline_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.045568, + "submission_ms": 0.045568, + "synchronized_e2e_ms": 0.074624 + }, + "cold_baseline_public_raw_existing_shape_hit": null, + "cold_baseline_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.198304, + "submission_ms": 0.198304, + "synchronized_e2e_ms": 0.21936 + }, + "cold_baseline_public_raw_shape_miss": { + "host_enqueue_ms": 7.880169, + "submission_ms": 7.880169, + "synchronized_e2e_ms": 7.908201 + }, + "cold_candidate_precomputed_first_pointer": { + "host_enqueue_ms": 0.099969, + "submission_ms": 0.099969, + "synchronized_e2e_ms": 0.117761 + }, + "cold_candidate_precomputed_fresh_pointer": { + "host_enqueue_ms": 0.04496, + "submission_ms": 0.04496, + "synchronized_e2e_ms": 0.060896 + }, + "cold_candidate_precomputed_prepare_first_pointer": { + "host_enqueue_ms": 1.245441, + "submission_ms": 1.245441, + "synchronized_e2e_ms": 1.271361 + }, + "cold_candidate_precomputed_prepare_fresh_pointer": { + "host_enqueue_ms": 1.092737, + "submission_ms": 1.092737, + "synchronized_e2e_ms": 1.116257 + }, + "cold_candidate_public_raw_existing_shape_hit": null, + "cold_candidate_public_raw_fresh_pointer_hit": { + "host_enqueue_ms": 0.074976, + "submission_ms": 0.074976, + "synchronized_e2e_ms": 0.103744 + }, + "cold_candidate_public_raw_shape_miss": { + "host_enqueue_ms": 2.39517, + "submission_ms": 2.39517, + "synchronized_e2e_ms": 2.420834 + }, + "cold_measurement_order": [ + "candidate", + "baseline" + ], + "cold_measurement_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "dtype": "bfloat16", + "evolution_flashlib_ms": 0.335967, + "evolution_kernel_ms": 0.179807, + "evolution_speedup": 1.8685, + "evolution_tflops": 17.9149, + "expected_route": "gap_pad_to_supported_seed_v1", + "fresh_pointer_rebind_verified": true, + "label": "post_d895_d48_b4_n8192_k1024_d48", + "measurement_order": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "measurement_order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "measurement_order_seed": "flashlib-kmeans-export-paired-replay-v1", + "measurement_session_id": "429dbdfd2f734f02819474b9385e8bb8", + "paired_cupti_session": { + "adaptive_probe_cupti_session_count": 4, + "adaptive_probes_reportable": false, + "adjacent_candidate_baseline_pairs": true, + "first_reportable_roles": [ + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", + "baseline_precomputed", + "candidate_precomputed", + "baseline_public_raw", + "candidate_public_raw", + "baseline_precomputed", + "candidate_precomputed", + "candidate_public_raw", + "baseline_public_raw", + "baseline_precomputed", + "candidate_precomputed" + ], + "interleaved": true, + "measurement_schedule_iteration_count": 12978, + "measurement_schedule_sha256": "410f77600189ecd8f6df959a46c7b52b90e94031c12487b6a68b4cbf317d2111", + "order_seed": "flashlib-kmeans-export-paired-replay-v1", + "precomputed_pair_count": 4083, + "public_pair_count": 2406, + "reportable_cupti_session_count": 1, + "role_sample_counts": { + "baseline_precomputed": 4083, + "baseline_public_raw": 2406, + "candidate_precomputed": 4083, + "candidate_public_raw": 2406 + }, + "same_cupti_session": true, + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "warmup_schedule_iteration_count": 2598 + }, + "paired_cupti_session_id": "429dbdfd2f734f02819474b9385e8bb8:post_d895_d48_b4_n8192_k1024_d48", + "precomputed_baseline_name": "triton_h200_07cf_precomputed", + "precomputed_gpu_speedup_vs_07cf": 0.9961038961038962, + "public_raw_baseline_name": "triton_h200_07cf_raw_adapter_v1", + "public_raw_e2e_speedup_vs_07cf_adapter": 1.9080601783744024, + "route_matches_expected": true, + "runtime_coverage": false, + "runtime_lifecycle_comparison": { + "amortized": [ + { + "after_init_synchronized_e2e_speedup": 3.2667258473732605, + "including_init_synchronized_e2e_speedup": 11.617177111352373, + "public_call_count": 1 + }, + { + "after_init_synchronized_e2e_speedup": 2.9346521090194284, + "including_init_synchronized_e2e_speedup": 11.43009478530128, + "public_call_count": 10 + }, + { + "after_init_synchronized_e2e_speedup": 2.2061313632197415, + "including_init_synchronized_e2e_speedup": 9.891741086372386, + "public_call_count": 100 + }, + { + "after_init_synchronized_e2e_speedup": 1.9448750073677699, + "including_init_synchronized_e2e_speedup": 4.960424811040331, + "public_call_count": 1000 + } + ], + "hot_gpu_span_speedup": 1.6468101460415063, + "hot_synchronized_e2e_speedup": 1.9080601783744024, + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency" + }, + "seed": 24802, + "selected_route": "gap_pad_to_supported_seed_v1", + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape": "post_d895_d48_b4_n8192_k1024_d48", + "source": "new_lowmid_gap", + "timing_backend": "cupti", + "triton_h200_07cf_config": { + "BLOCK_K": 64, + "BLOCK_N": 64, + "D_PAD": 64, + "kernel": "triton_h200_small_d", + "num_stages": 1, + "num_warps": 4 + } + } + ], + "runtime_lifecycle": { + "amortization_call_counts": [ + 1, + 10, + 100, + 1000 + ], + "baseline_api": "triton_h200_07cf_raw_adapter_v1.compute", + "baseline_has_explicit_init": true, + "baseline_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e", + "cache_policy": "synchronize_and_clear_after_each_completed_shape", + "candidate_api": "flashlib_cake_kmeans.init(...).compute", + "candidate_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion", + "cold_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "init_composition": "runtime_init_plus_shared_preprocess_support_each_standalone_lane", + "init_order_policy": "alternate_candidate_baseline_first_by_validation_shard_parity_then_shared_support", + "init_scope": "runtime_init_plus_standalone_shared_preprocess_support_once_per_validation_shard_process_device_operator", + "resident_multi_shape_cache_benchmarked": false, + "schema": "loom-public-runtime-lifecycle-v1", + "sessions": [ + { + "baseline_init": { + "host_enqueue_ms": 34.227907, + "submission_ms": 34.227907, + "synchronized_e2e_ms": 34.345443 + }, + "candidate_init": { + "host_enqueue_ms": 34.57754, + "submission_ms": 34.57754, + "synchronized_e2e_ms": 448.087792 + }, + "clear_count": { + "baseline": 31, + "candidate": 31 + }, + "final_baseline_cache_info": { + "hits": 0, + "misses": 0, + "size": 0 + }, + "final_candidate_cache_info": { + "hits": 0, + "max_cached_shapes": null, + "misses": 0, + "size": 0 + }, + "first_lookup_miss_count": { + "baseline": 31, + "candidate": 31 + }, + "fresh_pointer_hit_count": { + "baseline": 31, + "candidate": 31 + }, + "id": "a591a388bda74a2890fb7b17b61b2608", + "init_components": { + "attribution": "included_once_in_each_standalone_lane_init_total", + "baseline_runtime_init": { + "host_enqueue_ms": 0.043328, + "submission_ms": 0.043328, + "synchronized_e2e_ms": 0.133152 + }, + "candidate_runtime_init": { + "host_enqueue_ms": 0.392961, + "submission_ms": 0.392961, + "synchronized_e2e_ms": 413.875501 + }, + "shared_preprocess_compile": { + "host_enqueue_ms": 34.184579, + "submission_ms": 34.184579, + "synchronized_e2e_ms": 34.212291 + } + }, + "init_measurement_order": [ + "candidate", + "baseline", + "shared_preprocess_compile" + ], + "validation_shard": { + "count": 4, + "index": 0 + } + }, + { + "baseline_init": { + "host_enqueue_ms": 34.449444, + "submission_ms": 34.449444, + "synchronized_e2e_ms": 484.041494 + }, + "candidate_init": { + "host_enqueue_ms": 34.874083999999996, + "submission_ms": 34.874083999999996, + "synchronized_e2e_ms": 34.943076 + }, + "clear_count": { + "baseline": 31, + "candidate": 31 + }, + "final_baseline_cache_info": { + "hits": 0, + "misses": 0, + "size": 0 + }, + "final_candidate_cache_info": { + "hits": 0, + "max_cached_shapes": null, + "misses": 0, + "size": 0 + }, + "first_lookup_miss_count": { + "baseline": 31, + "candidate": 31 + }, + "fresh_pointer_hit_count": { + "baseline": 31, + "candidate": 31 + }, + "id": "41892d8f27794f7aba63e10273932712", + "init_components": { + "attribution": "included_once_in_each_standalone_lane_init_total", + "baseline_runtime_init": { + "host_enqueue_ms": 0.029344, + "submission_ms": 0.029344, + "synchronized_e2e_ms": 449.598162 + }, + "candidate_runtime_init": { + "host_enqueue_ms": 0.453984, + "submission_ms": 0.453984, + "synchronized_e2e_ms": 0.499744 + }, + "shared_preprocess_compile": { + "host_enqueue_ms": 34.4201, + "submission_ms": 34.4201, + "synchronized_e2e_ms": 34.443332 + } + }, + "init_measurement_order": [ + "baseline", + "candidate", + "shared_preprocess_compile" + ], + "validation_shard": { + "count": 4, + "index": 1 + } + }, + { + "baseline_init": { + "host_enqueue_ms": 35.551268, + "submission_ms": 35.551268, + "synchronized_e2e_ms": 35.631237 + }, + "candidate_init": { + "host_enqueue_ms": 35.83802, + "submission_ms": 35.83802, + "synchronized_e2e_ms": 484.100213 + }, + "clear_count": { + "baseline": 31, + "candidate": 31 + }, + "final_baseline_cache_info": { + "hits": 0, + "misses": 0, + "size": 0 + }, + "final_candidate_cache_info": { + "hits": 0, + "max_cached_shapes": null, + "misses": 0, + "size": 0 + }, + "first_lookup_miss_count": { + "baseline": 31, + "candidate": 31 + }, + "fresh_pointer_hit_count": { + "baseline": 31, + "candidate": 31 + }, + "id": "958ecd22358d4d5cb8ab46b2ff2cd831", + "init_components": { + "attribution": "included_once_in_each_standalone_lane_init_total", + "baseline_runtime_init": { + "host_enqueue_ms": 0.027136, + "submission_ms": 0.027136, + "synchronized_e2e_ms": 0.084128 + }, + "candidate_runtime_init": { + "host_enqueue_ms": 0.313888, + "submission_ms": 0.313888, + "synchronized_e2e_ms": 448.553104 + }, + "shared_preprocess_compile": { + "host_enqueue_ms": 35.524132, + "submission_ms": 35.524132, + "synchronized_e2e_ms": 35.547109 + } + }, + "init_measurement_order": [ + "candidate", + "baseline", + "shared_preprocess_compile" + ], + "validation_shard": { + "count": 4, + "index": 2 + } + }, + { + "baseline_init": { + "host_enqueue_ms": 36.934438, + "submission_ms": 36.934438, + "synchronized_e2e_ms": 455.11131900000004 + }, + "candidate_init": { + "host_enqueue_ms": 37.358022999999996, + "submission_ms": 37.358022999999996, + "synchronized_e2e_ms": 37.435623 + }, + "clear_count": { + "baseline": 31, + "candidate": 31 + }, + "final_baseline_cache_info": { + "hits": 0, + "misses": 0, + "size": 0 + }, + "final_candidate_cache_info": { + "hits": 0, + "max_cached_shapes": null, + "misses": 0, + "size": 0 + }, + "first_lookup_miss_count": { + "baseline": 31, + "candidate": 31 + }, + "fresh_pointer_hit_count": { + "baseline": 31, + "candidate": 31 + }, + "id": "429dbdfd2f734f02819474b9385e8bb8", + "init_components": { + "attribution": "included_once_in_each_standalone_lane_init_total", + "baseline_runtime_init": { + "host_enqueue_ms": 0.02192, + "submission_ms": 0.02192, + "synchronized_e2e_ms": 418.175089 + }, + "candidate_runtime_init": { + "host_enqueue_ms": 0.445505, + "submission_ms": 0.445505, + "synchronized_e2e_ms": 0.499393 + }, + "shared_preprocess_compile": { + "host_enqueue_ms": 36.912518, + "submission_ms": 36.912518, + "synchronized_e2e_ms": 36.93623 + } + }, + "init_measurement_order": [ + "baseline", + "candidate", + "shared_preprocess_compile" + ], + "validation_shard": { + "count": 4, + "index": 3 + } + } + ] + }, + "runtime_lifecycle_summary": { + "amortized_after_init_synchronized_e2e_speedup": { + "1": { + "geomean": 1.2309434323137933, + "max": 163.79859302023343, + "median": 2.7839934141126323, + "min": 0.08647202905817367, + "p90": 3.351566582003962, + "sample_count": 124 + }, + "10": { + "geomean": 1.2588822989268478, + "max": 137.34998767211528, + "median": 2.6667329239063666, + "min": 0.11096680845482784, + "p90": 3.090755401506125, + "sample_count": 124 + }, + "100": { + "geomean": 1.509344986064645, + "max": 53.5581614291322, + "median": 2.1681602849557455, + "min": 0.28495469134404755, + "p90": 2.989410840698895, + "sample_count": 124 + }, + "1000": { + "geomean": 1.9957356325578275, + "max": 9.526755432882533, + "median": 2.02714552137833, + "min": 1.0119821488655745, + "p90": 2.8424059645989446, + "sample_count": 124 + } + }, + "amortized_including_init_synchronized_e2e_speedup": { + "1": { + "geomean": 0.9449946730011419, + "max": 24.89092858456689, + "median": 2.2809542656641058, + "min": 0.07723037119783416, + "p90": 12.93634198448204, + "sample_count": 124 + }, + "10": { + "geomean": 0.9590724810364581, + "max": 24.52727808638291, + "median": 2.2785747135963885, + "min": 0.08144088410284854, + "p90": 12.756198569243546, + "sample_count": 124 + }, + "100": { + "geomean": 1.0709455274785604, + "max": 21.448324006518714, + "median": 2.256511456266482, + "min": 0.1092928565392167, + "p90": 10.92530869455313, + "sample_count": 124 + }, + "1000": { + "geomean": 1.4680428292344323, + "max": 10.378991521881256, + "median": 1.5029920226671636, + "min": 0.3509947998170528, + "p90": 6.045445059831923, + "sample_count": 124 + } + }, + "baseline_has_explicit_init": true, + "baseline_hot_public_synchronized_e2e_ms": { + "geomean": 0.20757016572840578, + "max": 0.5228809999999999, + "median": 0.18443199999999998, + "min": 0.16128, + "p90": 0.32694080000000014, + "sample_count": 124 + }, + "baseline_init_host_enqueue_ms": { + "geomean": 35.27461542656189, + "max": 36.934438, + "median": 35.000356, + "min": 34.227907, + "p90": 36.519487, + "sample_count": 4 + }, + "baseline_init_synchronized_e2e_ms": { + "geomean": 128.13715080146187, + "max": 484.041494, + "median": 245.37127800000002, + "min": 34.345443, + "p90": 475.36244150000005, + "sample_count": 4 + }, + "cache_policy": "synchronize_and_clear_after_each_completed_shape", + "candidate_first_compute_host_enqueue_ms": { + "geomean": 7.402417492079541, + "max": 83.722038, + "median": 3.0477155, + "min": 1.974882, + "p90": 70.15640230000001, + "sample_count": 124 + }, + "candidate_first_compute_synchronized_e2e_ms": { + "geomean": 7.475266319780226, + "max": 83.779798, + "median": 3.1303549999999998, + "min": 2.05613, + "p90": 70.21168870000001, + "sample_count": 124 + }, + "candidate_hot_compute_gpu_span_ms": { + "geomean": 0.03929959622925097, + "max": 0.389696, + "median": 0.038976, + "min": 0.009792, + "p90": 0.1054144, + "sample_count": 124 + }, + "candidate_hot_compute_host_enqueue_ms": { + "geomean": 0.04880204368462327, + "max": 0.06767999999999999, + "median": 0.047231999999999996, + "min": 0.040064, + "p90": 0.0595648, + "sample_count": 124 + }, + "candidate_hot_compute_synchronized_e2e_ms": { + "geomean": 0.09510640113954803, + "max": 0.438208, + "median": 0.08894425, + "min": 0.056224, + "p90": 0.1581216, + "sample_count": 124 + }, + "candidate_init_host_enqueue_ms": { + "geomean": 35.64562210040985, + "max": 37.358022999999996, + "median": 35.356052, + "min": 34.57754, + "p90": 36.902022099999996, + "sample_count": 4 + }, + "candidate_init_synchronized_e2e_ms": { + "geomean": 129.7884798704558, + "max": 484.100213, + "median": 242.7617075, + "min": 34.943076, + "p90": 473.29648670000006, + "sample_count": 4 + }, + "cold_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "hot_gpu_span_speedup": { + "geomean": 2.316300717543746, + "max": 6.406190408085431, + "median": 2.159249130548963, + "min": 0.9709337534898999, + "p90": 4.868219371609404, + "sample_count": 124 + }, + "hot_synchronized_e2e_speedup": { + "geomean": 2.1825046815076243, + "max": 3.8773006134969323, + "median": 2.1942486496321694, + "min": 1.093838998831605, + "p90": 2.842345086787142, + "sample_count": 124 + }, + "init_composition": "runtime_init_plus_shared_preprocess_support_each_standalone_lane", + "init_order_policy": "alternate_candidate_baseline_first_by_validation_shard_parity_then_shared_support", + "init_scope": "runtime_init_plus_standalone_shared_preprocess_support_once_per_validation_shard_process_device_operator", + "resident_multi_shape_cache_benchmarked": false, + "schema": "loom-public-runtime-lifecycle-summary-v1", + "shape_count": 124, + "validation_shard_count": 4 + }, + "runtime_workspace_lifecycle": "both_init_once_runtimes_synchronize_and_clear_after_each_completed_shape", + "schema": "flash-kmeans-export-dual-lane-v2", + "selected_row_count": 124, + "selected_unique_shape_count": 123, + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shapes": [ + { + "B": 1, + "D": 112, + "K": 4096, + "N": 1408, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.4908, + "evolution_flashlib_ms": 0.866559, + "evolution_kernel_ms": 0.196192, + "evolution_speedup": 4.4169, + "evolution_tflops": 6.5846, + "label": "adjacent_1d49_d112_highk_tail_b1_n1408_k4096_d112", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 1491121, + "source": "tail_divisibility" + }, + { + "B": 5, + "D": 128, + "K": 512, + "N": 7296, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 14.0693, + "evolution_flashlib_ms": 0.339855, + "evolution_kernel_ms": 0.152639, + "evolution_speedup": 2.2265, + "evolution_tflops": 31.3255, + "label": "adjacent_1d49_d128_forced_fallback_b5_n7296_k512_d128", + "route": "aligned_v10_fallback", + "runtime_coverage": false, + "seed": 1491281, + "source": "forced_fallback" + }, + { + "B": 5, + "D": 224, + "K": 768, + "N": 5632, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 20.8173, + "evolution_flashlib_ms": 0.465423, + "evolution_kernel_ms": 0.186336, + "evolution_speedup": 2.4978, + "evolution_tflops": 51.9966, + "label": "adjacent_1d49_d224_highn_overlap_b5_n5632_k768_d224", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 1492241, + "source": "guard_overlap" + }, + { + "B": 3, + "D": 288, + "K": 2048, + "N": 1152, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 4.0485, + "evolution_flashlib_ms": 1.007006, + "evolution_kernel_ms": 0.207744, + "evolution_speedup": 4.8473, + "evolution_tflops": 19.6245, + "label": "adjacent_1d49_d288_low_n_heldout_b3_n1152_k2048_d288", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 1492881, + "source": "heldout_neighborhood" + }, + { + "B": 2, + "D": 352, + "K": 8192, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 4.5385, + "evolution_flashlib_ms": 2.602461, + "evolution_kernel_ms": 0.302655, + "evolution_speedup": 8.5988, + "evolution_tflops": 39.0252, + "label": "adjacent_1d49_d352_request_highk_b2_n1024_k8192_d352", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 1493521, + "source": "request_specific" + }, + { + "B": 4, + "D": 416, + "K": 512, + "N": 3840, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 16.9546, + "evolution_flashlib_ms": 0.385919, + "evolution_kernel_ms": 0.184, + "evolution_speedup": 2.0974, + "evolution_tflops": 35.5604, + "label": "adjacent_1d49_d416_random_b4_n3840_k512_d416", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 1494161, + "source": "random_legal" + }, + { + "B": 2, + "D": 480, + "K": 256, + "N": 2816, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 4.1626, + "evolution_flashlib_ms": 0.332511, + "evolution_kernel_ms": 0.178015, + "evolution_speedup": 1.8679, + "evolution_tflops": 7.7753, + "label": "adjacent_1d49_d480_smallk_boundary_b2_n2816_k256_d480", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 1494801, + "source": "guard_boundary" + }, + { + "B": 4, + "D": 48, + "K": 256, + "N": 3968, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.7051, + "evolution_flashlib_ms": 0.228768, + "evolution_kernel_ms": 0.17192, + "evolution_speedup": 1.3307, + "evolution_tflops": 2.2689, + "label": "adjacent_1d49_d48_lowk_boundary_b4_n3968_k256_d48", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 1490481, + "source": "guard_boundary" + }, + { + "B": 2, + "D": 112, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.0686, + "evolution_flashlib_ms": 1.758429, + "evolution_kernel_ms": 0.295872, + "evolution_speedup": 5.9432, + "evolution_tflops": 6.3509, + "label": "adjacent_3328_d112_highk_low_n_b2_n512_k8192_d112", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 3328112, + "source": "heldout_neighborhood" + }, + { + "B": 8, + "D": 128, + "K": 512, + "N": 8192, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 21.7692, + "evolution_flashlib_ms": 0.394592, + "evolution_kernel_ms": 0.177472, + "evolution_speedup": 2.2234, + "evolution_tflops": 48.4016, + "label": "adjacent_3328_d128_exact_guard_k_neighbor_b8_n8192_k512_d128", + "route": "paired_large_v15", + "runtime_coverage": false, + "seed": 3328128, + "source": "guard_miss_fallback" + }, + { + "B": 3, + "D": 224, + "K": 512, + "N": 3840, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 8.3041, + "evolution_flashlib_ms": 0.318207, + "evolution_kernel_ms": 0.174895, + "evolution_speedup": 1.8194, + "evolution_tflops": 15.1086, + "label": "adjacent_3328_d224_tail_div_b3_n3840_k512_d224", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 3328224, + "source": "tail_divisibility" + }, + { + "B": 4, + "D": 288, + "K": 256, + "N": 8192, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 14.007, + "evolution_flashlib_ms": 0.344959, + "evolution_kernel_ms": 0.182432, + "evolution_speedup": 1.8909, + "evolution_tflops": 26.4857, + "label": "adjacent_3328_d288_guard_overlap_b4_n8192_k256_d288", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 3328288, + "source": "guard_overlap" + }, + { + "B": 3, + "D": 352, + "K": 768, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 5.7044, + "evolution_flashlib_ms": 0.582335, + "evolution_kernel_ms": 0.183872, + "evolution_speedup": 3.1671, + "evolution_tflops": 18.0663, + "label": "adjacent_3328_d352_random_legal_b3_n2048_k768_d352", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 3328352, + "source": "random_legal" + }, + { + "B": 1, + "D": 416, + "K": 8192, + "N": 384, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.063, + "evolution_flashlib_ms": 41.538691, + "evolution_kernel_ms": 0.276031, + "evolution_speedup": 150.4856, + "evolution_tflops": 9.4817, + "label": "adjacent_3328_d416_highk_low_n_b1_n384_k8192_d416", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 3328416, + "source": "request_specific" + }, + { + "B": 1, + "D": 480, + "K": 256, + "N": 128, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0644, + "evolution_flashlib_ms": 0.488575, + "evolution_kernel_ms": 0.177056, + "evolution_speedup": 2.7594, + "evolution_tflops": 0.1777, + "label": "adjacent_3328_d480_min_boundary_b1_n128_k256_d480", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 3328480, + "source": "guard_boundary" + }, + { + "B": 1, + "D": 48, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0145, + "evolution_flashlib_ms": 0.434063, + "evolution_kernel_ms": 0.168656, + "evolution_speedup": 2.5737, + "evolution_tflops": 0.0373, + "label": "adjacent_3328_d48_small_boundary_b1_n256_k256_d48", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 3328481, + "source": "guard_boundary" + }, + { + "B": 5, + "D": 112, + "K": 512, + "N": 2176, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 3.6867, + "evolution_flashlib_ms": 0.338463, + "evolution_kernel_ms": 0.171519, + "evolution_speedup": 1.9733, + "evolution_tflops": 7.275, + "label": "adjacent_5600_d112_odd_tile_tail_b5_n2176_k512_d112", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 5601121, + "source": "tail_divisibility" + }, + { + "B": 8, + "D": 128, + "K": 256, + "N": 8320, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 14.6812, + "evolution_flashlib_ms": 0.297119, + "evolution_kernel_ms": 0.155872, + "evolution_speedup": 1.9062, + "evolution_tflops": 27.985, + "label": "adjacent_5600_d128_fallback_neighbor_b8_n8320_k256_d128", + "route": "aligned_v10_fallback", + "runtime_coverage": false, + "seed": 5601281, + "source": "guard_miss_fallback" + }, + { + "B": 3, + "D": 224, + "K": 512, + "N": 3072, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 6.148, + "evolution_flashlib_ms": 0.343839, + "evolution_kernel_ms": 0.174928, + "evolution_speedup": 1.9656, + "evolution_tflops": 12.0846, + "label": "adjacent_5600_d224_k512_overlap_b3_n3072_k512_d224", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 5602241, + "source": "guard_overlap" + }, + { + "B": 2, + "D": 288, + "K": 2048, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.4751, + "evolution_flashlib_ms": 0.976096, + "evolution_kernel_ms": 0.206512, + "evolution_speedup": 4.7266, + "evolution_tflops": 11.6987, + "label": "adjacent_5600_d288_mid_highk_b2_n1024_k2048_d288", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 5602881, + "source": "heldout_neighborhood" + }, + { + "B": 2, + "D": 352, + "K": 4096, + "N": 768, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.7323, + "evolution_flashlib_ms": 1.621022, + "evolution_kernel_ms": 1.051615, + "evolution_speedup": 1.5415, + "evolution_tflops": 4.2118, + "label": "adjacent_5600_d352_splitk_edge_b2_n768_k4096_d352", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 5603521, + "source": "request_specific" + }, + { + "B": 4, + "D": 416, + "K": 256, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 4.7141, + "evolution_flashlib_ms": 0.370128, + "evolution_kernel_ms": 0.175743, + "evolution_speedup": 2.1061, + "evolution_tflops": 9.9283, + "label": "adjacent_5600_d416_smallk_boundary_b4_n2048_k256_d416", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 5604161, + "source": "guard_boundary" + }, + { + "B": 1, + "D": 480, + "K": 1024, + "N": 1536, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.1001, + "evolution_flashlib_ms": 0.718975, + "evolution_kernel_ms": 0.199488, + "evolution_speedup": 3.6041, + "evolution_tflops": 7.5691, + "label": "adjacent_5600_d480_random_b1_n1536_k1024_d480", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 5604801, + "source": "random_legal" + }, + { + "B": 3, + "D": 48, + "K": 256, + "N": 1536, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.4029, + "evolution_flashlib_ms": 0.281087, + "evolution_kernel_ms": 0.169183, + "evolution_speedup": 1.6614, + "evolution_tflops": 0.6694, + "label": "adjacent_5600_d48_low_n_boundary_b3_n1536_k256_d48", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 5600481, + "source": "guard_boundary" + }, + { + "B": 5, + "D": 112, + "K": 512, + "N": 2944, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 5.7016, + "evolution_flashlib_ms": 0.296095, + "evolution_kernel_ms": 0.173023, + "evolution_speedup": 1.7113, + "evolution_tflops": 9.7571, + "label": "adjacent_68cf_d112_tail_b5_n2944_k512_d112", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 6811201, + "source": "tail_divisibility" + }, + { + "B": 7, + "D": 128, + "K": 512, + "N": 6016, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 14.2803, + "evolution_flashlib_ms": 0.386527, + "evolution_kernel_ms": 0.156832, + "evolution_speedup": 2.4646, + "evolution_tflops": 35.195, + "label": "adjacent_68cf_d128_forced_fallback_b7_n6016_k512_d128", + "route": "aligned_v10_fallback", + "runtime_coverage": false, + "seed": 6812801, + "source": "forced_fallback" + }, + { + "B": 3, + "D": 224, + "K": 1024, + "N": 5120, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 14.4926, + "evolution_flashlib_ms": 0.486208, + "evolution_kernel_ms": 0.180032, + "evolution_speedup": 2.7007, + "evolution_tflops": 39.1399, + "label": "adjacent_68cf_d224_overlap_b3_n5120_k1024_d224", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 6822401, + "source": "guard_overlap" + }, + { + "B": 2, + "D": 288, + "K": 512, + "N": 1920, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.6899, + "evolution_flashlib_ms": 0.421007, + "evolution_kernel_ms": 0.177824, + "evolution_speedup": 2.3676, + "evolution_tflops": 6.3684, + "label": "adjacent_68cf_d288_boundary_b2_n1920_k512_d288", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 6828801, + "source": "guard_boundary" + }, + { + "B": 3, + "D": 352, + "K": 768, + "N": 2816, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 9.1314, + "evolution_flashlib_ms": 0.500207, + "evolution_kernel_ms": 0.18768, + "evolution_speedup": 2.6652, + "evolution_tflops": 24.3372, + "label": "adjacent_68cf_d352_tail_b3_n2816_k768_d352", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 6835201, + "source": "tail_divisibility" + }, + { + "B": 2, + "D": 416, + "K": 1024, + "N": 2304, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 6.4021, + "evolution_flashlib_ms": 0.613215, + "evolution_kernel_ms": 0.196895, + "evolution_speedup": 3.1144, + "evolution_tflops": 19.9389, + "label": "adjacent_68cf_d416_overlap_b2_n2304_k1024_d416", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 6841601, + "source": "guard_overlap" + }, + { + "B": 4, + "D": 480, + "K": 512, + "N": 1664, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 6.0602, + "evolution_flashlib_ms": 0.539839, + "evolution_kernel_ms": 0.18536, + "evolution_speedup": 2.9124, + "evolution_tflops": 17.6497, + "label": "adjacent_68cf_d480_boundary_b4_n1664_k512_d480", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 6848001, + "source": "guard_boundary" + }, + { + "B": 4, + "D": 48, + "K": 512, + "N": 2304, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.5408, + "evolution_flashlib_ms": 0.293984, + "evolution_kernel_ms": 0.170608, + "evolution_speedup": 1.7232, + "evolution_tflops": 2.6551, + "label": "adjacent_68cf_d48_boundary_b4_n2304_k512_d48", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 6804801, + "source": "guard_boundary" + }, + { + "B": 1, + "D": 112, + "K": 4096, + "N": 384, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.2782, + "evolution_flashlib_ms": 1.266206, + "evolution_kernel_ms": 0.195296, + "evolution_speedup": 6.4835, + "evolution_tflops": 1.804, + "label": "adjacent_8f09_d112_small_highk_b1_n384_k4096_d112", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 8091121, + "source": "heldout_neighborhood" + }, + { + "B": 4, + "D": 128, + "K": 512, + "N": 4480, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 7.2152, + "evolution_flashlib_ms": 0.325535, + "evolution_kernel_ms": 0.145856, + "evolution_speedup": 2.2319, + "evolution_tflops": 16.1036, + "label": "adjacent_8f09_d128_fallback_neighbor_b4_n4480_k512_d128", + "route": "aligned_v10_fallback", + "runtime_coverage": false, + "seed": 8091281, + "source": "guard_miss_fallback" + }, + { + "B": 4, + "D": 224, + "K": 256, + "N": 1536, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.2208, + "evolution_flashlib_ms": 0.317296, + "evolution_kernel_ms": 0.173664, + "evolution_speedup": 1.8271, + "evolution_tflops": 4.0575, + "label": "adjacent_8f09_d224_low_n_boundary_b4_n1536_k256_d224", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 8092241, + "source": "guard_boundary" + }, + { + "B": 1, + "D": 288, + "K": 768, + "N": 2560, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.3579, + "evolution_flashlib_ms": 0.480287, + "evolution_kernel_ms": 0.18432, + "evolution_speedup": 2.6057, + "evolution_tflops": 6.144, + "label": "adjacent_8f09_d288_k768_overlap_b1_n2560_k768_d288", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 8092881, + "source": "guard_overlap" + }, + { + "B": 4, + "D": 352, + "K": 256, + "N": 4096, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 9.629, + "evolution_flashlib_ms": 0.306655, + "evolution_kernel_ms": 0.176128, + "evolution_speedup": 1.7411, + "evolution_tflops": 16.765, + "label": "adjacent_8f09_d352_smallk_large_n_b4_n4096_k256_d352", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 8093521, + "source": "random_legal" + }, + { + "B": 3, + "D": 416, + "K": 768, + "N": 3456, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 13.4164, + "evolution_flashlib_ms": 0.493791, + "evolution_kernel_ms": 0.189088, + "evolution_speedup": 2.6114, + "evolution_tflops": 35.0361, + "label": "adjacent_8f09_d416_midk_tail_b3_n3456_k768_d416", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 8094161, + "source": "tail_divisibility" + }, + { + "B": 2, + "D": 480, + "K": 4096, + "N": 640, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.5174, + "evolution_flashlib_ms": 1.999325, + "evolution_kernel_ms": 0.284192, + "evolution_speedup": 7.0351, + "evolution_tflops": 17.7104, + "label": "adjacent_8f09_d480_highk_low_n_b2_n640_k4096_d480", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 8094801, + "source": "request_specific" + }, + { + "B": 2, + "D": 48, + "K": 512, + "N": 1792, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.556, + "evolution_flashlib_ms": 0.316848, + "evolution_kernel_ms": 0.169375, + "evolution_speedup": 1.8707, + "evolution_tflops": 1.0401, + "label": "adjacent_8f09_d48_medium_tail_b2_n1792_k512_d48", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 8090481, + "source": "tail_divisibility" + }, + { + "B": 3, + "D": 112, + "K": 768, + "N": 3840, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 5.9472, + "evolution_flashlib_ms": 0.333232, + "evolution_kernel_ms": 0.194592, + "evolution_speedup": 1.7125, + "evolution_tflops": 10.1844, + "label": "adjacent_9ca0_d112_tail_div_b3_n3840_k768_d112", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 921125, + "source": "tail_divisibility" + }, + { + "B": 8, + "D": 128, + "K": 256, + "N": 8064, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 14.5156, + "evolution_flashlib_ms": 0.291264, + "evolution_kernel_ms": 0.155168, + "evolution_speedup": 1.8771, + "evolution_tflops": 27.247, + "label": "adjacent_9ca0_d128_exact_guard_neighbor_b8_n8064_k256_d128", + "route": "aligned_v10_fallback", + "runtime_coverage": false, + "seed": 912806, + "source": "guard_miss_fallback" + }, + { + "B": 2, + "D": 224, + "K": 256, + "N": 4096, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 3.6395, + "evolution_flashlib_ms": 0.258143, + "evolution_kernel_ms": 0.1712, + "evolution_speedup": 1.5078, + "evolution_tflops": 5.4879, + "label": "adjacent_9ca0_d224_guard_overlap_b2_n4096_k256_d224", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 922405, + "source": "guard_overlap" + }, + { + "B": 1, + "D": 288, + "K": 4096, + "N": 384, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0159, + "evolution_flashlib_ms": 56.807529, + "evolution_kernel_ms": 0.243872, + "evolution_speedup": 232.9399, + "evolution_tflops": 3.7149, + "label": "adjacent_9ca0_d288_highk_low_n_b1_n384_k4096_d288", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 928805, + "source": "heldout_neighborhood" + }, + { + "B": 1, + "D": 352, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0601, + "evolution_flashlib_ms": 49.145822, + "evolution_kernel_ms": 0.275519, + "evolution_speedup": 178.3754, + "evolution_tflops": 10.7172, + "label": "adjacent_9ca0_d352_splitk_boundary_b1_n512_k8192_d352", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 935205, + "source": "request_specific" + }, + { + "B": 1, + "D": 416, + "K": 4096, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0848, + "evolution_flashlib_ms": 41.154087, + "evolution_kernel_ms": 0.274271, + "evolution_speedup": 150.049, + "evolution_tflops": 12.7234, + "label": "adjacent_9ca0_d416_splitk_min_b1_n1024_k4096_d416", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 941605, + "source": "guard_boundary" + }, + { + "B": 2, + "D": 480, + "K": 512, + "N": 4096, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 8.9917, + "evolution_flashlib_ms": 0.447807, + "evolution_kernel_ms": 0.185087, + "evolution_speedup": 2.4194, + "evolution_tflops": 21.7548, + "label": "adjacent_9ca0_d480_splitd_large_n_b2_n4096_k512_d480", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 948005, + "source": "random_legal" + }, + { + "B": 6, + "D": 48, + "K": 512, + "N": 12288, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 11.385, + "evolution_flashlib_ms": 0.318304, + "evolution_kernel_ms": 0.182336, + "evolution_speedup": 1.7457, + "evolution_tflops": 19.8747, + "label": "adjacent_9ca0_d48_guard_boundary_b6_n12288_k512_d48", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 924805, + "source": "guard_boundary" + }, + { + "B": 4, + "D": 112, + "K": 256, + "N": 3456, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 3.4264, + "evolution_flashlib_ms": 0.231359, + "evolution_kernel_ms": 0.168912, + "evolution_speedup": 1.3697, + "evolution_tflops": 4.6931, + "label": "adjacent_a2f8_d112_lowk_tail_b4_n3456_k256_d112", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 1028112, + "source": "tail_divisibility" + }, + { + "B": 6, + "D": 128, + "K": 512, + "N": 8576, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 19.2761, + "evolution_flashlib_ms": 0.349887, + "evolution_kernel_ms": 0.158688, + "evolution_speedup": 2.2049, + "evolution_tflops": 42.5013, + "label": "adjacent_a2f8_d128_odd_fallback_b6_n8576_k512_d128", + "route": "aligned_v10_fallback", + "runtime_coverage": false, + "seed": 1028128, + "source": "guard_miss_fallback" + }, + { + "B": 2, + "D": 224, + "K": 768, + "N": 6144, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 9.3487, + "evolution_flashlib_ms": 0.45224, + "evolution_kernel_ms": 0.176256, + "evolution_speedup": 2.5658, + "evolution_tflops": 23.987, + "label": "adjacent_a2f8_d224_highn_overlap_b2_n6144_k768_d224", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 1028224, + "source": "guard_overlap" + }, + { + "B": 1, + "D": 288, + "K": 4096, + "N": 640, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0262, + "evolution_flashlib_ms": 57.6586, + "evolution_kernel_ms": 0.24416, + "evolution_speedup": 236.1509, + "evolution_tflops": 6.1843, + "label": "adjacent_a2f8_d288_splitk_probe_b1_n640_k4096_d288", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 1028288, + "source": "heldout_neighborhood" + }, + { + "B": 4, + "D": 352, + "K": 256, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.5395, + "evolution_flashlib_ms": 0.479503, + "evolution_kernel_ms": 0.176096, + "evolution_speedup": 2.723, + "evolution_tflops": 4.192, + "label": "adjacent_a2f8_d352_smallk_boundary_b4_n1024_k256_d352", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 1028352, + "source": "guard_boundary" + }, + { + "B": 2, + "D": 416, + "K": 768, + "N": 2560, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 6.6441, + "evolution_flashlib_ms": 0.492399, + "evolution_kernel_ms": 0.189631, + "evolution_speedup": 2.5966, + "evolution_tflops": 17.2522, + "label": "adjacent_a2f8_d416_random_b2_n2560_k768_d416", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 1028416, + "source": "random_legal" + }, + { + "B": 1, + "D": 480, + "K": 4096, + "N": 896, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1098, + "evolution_flashlib_ms": 32.096809, + "evolution_kernel_ms": 0.274432, + "evolution_speedup": 116.9572, + "evolution_tflops": 12.8382, + "label": "adjacent_a2f8_d480_highk_request_b1_n896_k4096_d480", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 1028480, + "source": "request_specific" + }, + { + "B": 2, + "D": 48, + "K": 1024, + "N": 768, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.3038, + "evolution_flashlib_ms": 0.497055, + "evolution_kernel_ms": 0.173184, + "evolution_speedup": 2.8701, + "evolution_tflops": 0.8719, + "label": "adjacent_a2f8_d48_k1024_boundary_b2_n768_k1024_d48", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 1028481, + "source": "guard_boundary" + }, + { + "B": 2, + "D": 112, + "K": 768, + "N": 3200, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 3.4225, + "evolution_flashlib_ms": 0.321696, + "evolution_kernel_ms": 0.173088, + "evolution_speedup": 1.8586, + "evolution_tflops": 6.361, + "label": "adjacent_c44f_d112_odd_tail_b2_n3200_k768_d112", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 4411201, + "source": "tail_divisibility" + }, + { + "B": 6, + "D": 128, + "K": 512, + "N": 6272, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 14.7956, + "evolution_flashlib_ms": 0.333376, + "evolution_kernel_ms": 0.152256, + "evolution_speedup": 2.1896, + "evolution_tflops": 32.3961, + "label": "adjacent_c44f_d128_forced_fallback_b6_n6272_k512_d128", + "route": "aligned_v10_fallback", + "runtime_coverage": false, + "seed": 4412801, + "source": "forced_fallback" + }, + { + "B": 4, + "D": 224, + "K": 512, + "N": 4480, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 10.201, + "evolution_flashlib_ms": 0.402944, + "evolution_kernel_ms": 0.173008, + "evolution_speedup": 2.329, + "evolution_tflops": 23.7585, + "label": "adjacent_c44f_d224_overlap_b4_n4480_k512_d224", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 4422401, + "source": "guard_overlap" + }, + { + "B": 2, + "D": 288, + "K": 4096, + "N": 768, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.2892, + "evolution_flashlib_ms": 1.583007, + "evolution_kernel_ms": 0.246496, + "evolution_speedup": 6.422, + "evolution_tflops": 14.7016, + "label": "adjacent_c44f_d288_highk_heldout_b2_n768_k4096_d288", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 4428801, + "source": "heldout_neighborhood" + }, + { + "B": 1, + "D": 352, + "K": 768, + "N": 3328, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 3.6993, + "evolution_flashlib_ms": 0.4864, + "evolution_kernel_ms": 0.187632, + "evolution_speedup": 2.5923, + "evolution_tflops": 9.5898, + "label": "adjacent_c44f_d352_random_b1_n3328_k768_d352", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 4435201, + "source": "random_legal" + }, + { + "B": 3, + "D": 416, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.7386, + "evolution_flashlib_ms": 3.822732, + "evolution_kernel_ms": 0.303423, + "evolution_speedup": 12.5987, + "evolution_tflops": 34.5029, + "label": "adjacent_c44f_d416_highk_request_b3_n512_k8192_d416", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 4441601, + "source": "request_specific" + }, + { + "B": 5, + "D": 480, + "K": 512, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 10.031, + "evolution_flashlib_ms": 0.501759, + "evolution_kernel_ms": 0.185136, + "evolution_speedup": 2.7102, + "evolution_tflops": 27.1863, + "label": "adjacent_c44f_d480_boundary_b5_n2048_k512_d480", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 4448001, + "source": "guard_boundary" + }, + { + "B": 3, + "D": 48, + "K": 768, + "N": 2688, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.9182, + "evolution_flashlib_ms": 0.309952, + "evolution_kernel_ms": 0.174383, + "evolution_speedup": 1.7774, + "evolution_tflops": 3.4094, + "label": "adjacent_c44f_d48_midk_boundary_b3_n2688_k768_d48", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 4404801, + "source": "guard_boundary" + }, + { + "B": 3, + "D": 112, + "K": 8192, + "N": 768, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.7846, + "evolution_flashlib_ms": 1.518318, + "evolution_kernel_ms": 0.287872, + "evolution_speedup": 5.2743, + "evolution_tflops": 14.6866, + "label": "adjacent_d9d5_d112_highk_request_b3_n768_k8192_d112", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 9511201, + "source": "request_specific" + }, + { + "B": 5, + "D": 128, + "K": 1024, + "N": 6016, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 17.1342, + "evolution_flashlib_ms": 0.460207, + "evolution_kernel_ms": 0.159104, + "evolution_speedup": 2.8925, + "evolution_tflops": 49.5606, + "label": "adjacent_d9d5_d128_guard_miss_fallback_b5_n6016_k1024_d128", + "route": "aligned_v10_fallback", + "runtime_coverage": false, + "seed": 9512801, + "source": "guard_miss_fallback" + }, + { + "B": 2, + "D": 224, + "K": 1280, + "N": 2944, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 6.9153, + "evolution_flashlib_ms": 0.488255, + "evolution_kernel_ms": 0.180384, + "evolution_speedup": 2.7068, + "evolution_tflops": 18.7179, + "label": "adjacent_d9d5_d224_random_midk_b2_n2944_k1280_d224", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 9522401, + "source": "random_legal" + }, + { + "B": 1, + "D": 288, + "K": 4096, + "N": 896, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0365, + "evolution_flashlib_ms": 57.939113, + "evolution_kernel_ms": 0.2448, + "evolution_speedup": 236.6794, + "evolution_tflops": 8.6353, + "label": "adjacent_d9d5_d288_heldout_low_n_b1_n896_k4096_d288", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 9528801, + "source": "heldout_neighborhood" + }, + { + "B": 5, + "D": 352, + "K": 768, + "N": 2304, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 10.6911, + "evolution_flashlib_ms": 0.582592, + "evolution_kernel_ms": 0.186944, + "evolution_speedup": 3.1164, + "evolution_tflops": 33.3177, + "label": "adjacent_d9d5_d352_random_b5_n2304_k768_d352", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 9535201, + "source": "random_legal" + }, + { + "B": 2, + "D": 416, + "K": 8192, + "N": 640, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.6807, + "evolution_flashlib_ms": 3.254428, + "evolution_kernel_ms": 0.298432, + "evolution_speedup": 10.9051, + "evolution_tflops": 29.2333, + "label": "adjacent_d9d5_d416_highk_request_b2_n640_k8192_d416", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 9541601, + "source": "request_specific" + }, + { + "B": 3, + "D": 480, + "K": 4096, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 6.283, + "evolution_flashlib_ms": 1.92259, + "evolution_kernel_ms": 0.297632, + "evolution_speedup": 6.4596, + "evolution_tflops": 40.5857, + "label": "adjacent_d9d5_d480_heldout_highk_b3_n1024_k4096_d480", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 9548001, + "source": "heldout_neighborhood" + }, + { + "B": 2, + "D": 48, + "K": 4096, + "N": 640, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.5494, + "evolution_flashlib_ms": 0.916111, + "evolution_kernel_ms": 0.19392, + "evolution_speedup": 4.7242, + "evolution_tflops": 2.5955, + "label": "adjacent_d9d5_d48_highk_heldout_b2_n640_k4096_d48", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 9504801, + "source": "heldout_neighborhood" + }, + { + "B": 4, + "D": 112, + "K": 1024, + "N": 3712, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 8.805, + "evolution_flashlib_ms": 0.386799, + "evolution_kernel_ms": 0.176608, + "evolution_speedup": 2.1902, + "evolution_tflops": 19.2844, + "label": "adjacent_ef00_d112_odd_tail_b4_n3712_k1024_d112", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 9001121, + "source": "tail_divisibility" + }, + { + "B": 4, + "D": 128, + "K": 512, + "N": 7552, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 12.6084, + "evolution_flashlib_ms": 0.314032, + "evolution_kernel_ms": 0.153023, + "evolution_speedup": 2.0522, + "evolution_tflops": 25.8747, + "label": "adjacent_ef00_d128_forced_fallback_b4_n7552_k512_d128", + "route": "aligned_v10_fallback", + "runtime_coverage": false, + "seed": 9001281, + "source": "forced_fallback" + }, + { + "B": 2, + "D": 224, + "K": 512, + "N": 5376, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 6.6916, + "evolution_flashlib_ms": 0.36856, + "evolution_kernel_ms": 0.177808, + "evolution_speedup": 2.0728, + "evolution_tflops": 13.8703, + "label": "adjacent_ef00_d224_midn_overlap_b2_n5376_k512_d224", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 9002241, + "source": "guard_overlap" + }, + { + "B": 1, + "D": 288, + "K": 4096, + "N": 1664, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0667, + "evolution_flashlib_ms": 58.88955, + "evolution_kernel_ms": 0.247408, + "evolution_speedup": 238.0261, + "evolution_tflops": 15.868, + "label": "adjacent_ef00_d288_highk_heldout_b1_n1664_k4096_d288", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 9002881, + "source": "heldout_neighborhood" + }, + { + "B": 3, + "D": 352, + "K": 8192, + "N": 896, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 4.9423, + "evolution_flashlib_ms": 3.136652, + "evolution_kernel_ms": 0.316191, + "evolution_speedup": 9.9201, + "evolution_tflops": 49.0278, + "label": "adjacent_ef00_d352_request_low_n_b3_n896_k8192_d352", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 9003521, + "source": "request_specific" + }, + { + "B": 1, + "D": 416, + "K": 768, + "N": 2176, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.6268, + "evolution_flashlib_ms": 0.529311, + "evolution_kernel_ms": 0.191648, + "evolution_speedup": 2.7619, + "evolution_tflops": 7.255, + "label": "adjacent_ef00_d416_random_midk_b1_n2176_k768_d416", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 9004161, + "source": "random_legal" + }, + { + "B": 3, + "D": 480, + "K": 256, + "N": 3200, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 7.059, + "evolution_flashlib_ms": 0.334224, + "evolution_kernel_ms": 0.18192, + "evolution_speedup": 1.8372, + "evolution_tflops": 12.9689, + "label": "adjacent_ef00_d480_lowk_boundary_b3_n3200_k256_d480", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 9004801, + "source": "guard_boundary" + }, + { + "B": 3, + "D": 48, + "K": 512, + "N": 3456, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.9767, + "evolution_flashlib_ms": 0.257807, + "evolution_kernel_ms": 0.172495, + "evolution_speedup": 1.4946, + "evolution_tflops": 2.9543, + "label": "adjacent_ef00_d48_midn_boundary_b3_n3456_k512_d48", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 9000481, + "source": "guard_boundary" + }, + { + "B": 1, + "D": 112, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0353, + "evolution_flashlib_ms": 0.416031, + "evolution_kernel_ms": 0.166575, + "evolution_speedup": 2.4976, + "evolution_tflops": 0.0881, + "label": "post_d895_d112_b1_n256_k256_d112", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 211204, + "source": "lowmid_small_boundary" + }, + { + "B": 1, + "D": 112, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.5425, + "evolution_flashlib_ms": 1.731693, + "evolution_kernel_ms": 0.294304, + "evolution_speedup": 5.884, + "evolution_tflops": 3.1924, + "label": "post_d895_d112_b1_n512_k8192_d112", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 211203, + "source": "high_k_low_n" + }, + { + "B": 2, + "D": 112, + "K": 512, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.5026, + "evolution_flashlib_ms": 0.312623, + "evolution_kernel_ms": 0.188, + "evolution_speedup": 1.6629, + "evolution_tflops": 2.4987, + "label": "post_d895_d112_b2_n2048_k512_d112", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 211201, + "source": "new_lowmid_gap" + }, + { + "B": 4, + "D": 112, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 18.1068, + "evolution_flashlib_ms": 0.415103, + "evolution_kernel_ms": 0.1956, + "evolution_speedup": 2.1222, + "evolution_tflops": 38.4264, + "label": "post_d895_d112_b4_n8192_k1024_d112", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 211202, + "source": "new_lowmid_gap" + }, + { + "B": 3, + "D": 128, + "K": 256, + "N": 1920, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.3402, + "evolution_flashlib_ms": 0.281664, + "evolution_kernel_ms": 0.1436, + "evolution_speedup": 1.9614, + "evolution_tflops": 2.6287, + "label": "post_d895_d128_fallback_b3_n1920_k256_d128", + "route": "aligned_v10_fallback", + "runtime_coverage": false, + "seed": 212803, + "source": "near_floor_fallback_d128" + }, + { + "B": 5, + "D": 128, + "K": 512, + "N": 2176, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 4.6116, + "evolution_flashlib_ms": 0.309232, + "evolution_kernel_ms": 0.14432, + "evolution_speedup": 2.1427, + "evolution_tflops": 9.8813, + "label": "post_d895_d128_fallback_b5_n2176_k512_d128", + "route": "aligned_v10_fallback", + "runtime_coverage": false, + "seed": 212804, + "source": "near_floor_fallback_d128" + }, + { + "B": 7, + "D": 128, + "K": 1024, + "N": 2432, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 11.0028, + "evolution_flashlib_ms": 0.4056, + "evolution_kernel_ms": 0.147936, + "evolution_speedup": 2.7417, + "evolution_tflops": 30.1667, + "label": "post_d895_d128_fallback_b7_n2432_k1024_d128", + "route": "aligned_v10_fallback", + "runtime_coverage": false, + "seed": 212805, + "source": "near_floor_fallback_d128" + }, + { + "B": 2, + "D": 128, + "K": 256, + "N": 262144, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 51.8017, + "evolution_flashlib_ms": 0.663294, + "evolution_kernel_ms": 0.24464, + "evolution_speedup": 2.7113, + "evolution_tflops": 140.4505, + "label": "post_d895_d128_paired_b2_n262144_k256_d128", + "route": "paired_large_v15", + "runtime_coverage": false, + "seed": 212802, + "source": "near_floor_paired_d128" + }, + { + "B": 8, + "D": 128, + "K": 256, + "N": 8192, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 14.6751, + "evolution_flashlib_ms": 0.292671, + "evolution_kernel_ms": 0.15424, + "evolution_speedup": 1.8975, + "evolution_tflops": 27.846, + "label": "post_d895_d128_paired_b8_n8192_k256_d128", + "route": "d128_even_near_floor_v10_repair", + "runtime_coverage": false, + "seed": 212801, + "source": "near_floor_paired_d128" + }, + { + "B": 1, + "D": 144, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0351, + "evolution_flashlib_ms": 0.537279, + "evolution_kernel_ms": 0.167728, + "evolution_speedup": 3.2033, + "evolution_tflops": 0.1125, + "label": "post_d895_d144_b1_n256_k256_d144", + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "runtime_coverage": false, + "seed": 214401, + "source": "tailpad_boundary" + }, + { + "B": 1, + "D": 144, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.6035, + "evolution_flashlib_ms": 2.001693, + "evolution_kernel_ms": 0.302528, + "evolution_speedup": 6.6166, + "evolution_tflops": 3.9929, + "label": "post_d895_d144_b1_n512_k8192_d144", + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "runtime_coverage": false, + "seed": 214404, + "source": "high_k_low_n" + }, + { + "B": 2, + "D": 144, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.5941, + "evolution_flashlib_ms": 0.465663, + "evolution_kernel_ms": 0.192896, + "evolution_speedup": 2.4141, + "evolution_tflops": 6.2622, + "label": "post_d895_d144_b2_n2048_k1024_d144", + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "runtime_coverage": false, + "seed": 214402, + "source": "tailpad_boundary" + }, + { + "B": 4, + "D": 144, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 20.3965, + "evolution_flashlib_ms": 0.473791, + "evolution_kernel_ms": 0.195104, + "evolution_speedup": 2.4284, + "evolution_tflops": 49.5309, + "label": "post_d895_d144_b4_n8192_k1024_d144", + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "runtime_coverage": false, + "seed": 214403, + "source": "tailpad_boundary" + }, + { + "B": 4, + "D": 16, + "K": 1024, + "N": 32768, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 12.993, + "evolution_flashlib_ms": 0.330559, + "evolution_kernel_ms": 0.216591, + "evolution_speedup": 1.5262, + "evolution_tflops": 19.8298, + "label": "post_d895_d16_b4_n32768_k1024_d16", + "route": "microdim_hybrid_9c0d_v1", + "runtime_coverage": false, + "seed": 21602, + "source": "near_floor_microdim" + }, + { + "B": 8, + "D": 16, + "K": 512, + "N": 65536, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 19.3119, + "evolution_flashlib_ms": 0.444799, + "evolution_kernel_ms": 0.295359, + "evolution_speedup": 1.506, + "evolution_tflops": 29.083, + "label": "post_d895_d16_b8_n65536_k512_d16", + "route": "microdim_hybrid_9c0d_v1", + "runtime_coverage": false, + "seed": 21601, + "source": "near_floor_microdim" + }, + { + "B": 1, + "D": 176, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0453, + "evolution_flashlib_ms": 0.509327, + "evolution_kernel_ms": 0.169792, + "evolution_speedup": 2.9997, + "evolution_tflops": 0.1359, + "label": "post_d895_d176_b1_n256_k256_d176", + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "runtime_coverage": false, + "seed": 217601, + "source": "tailpad_boundary" + }, + { + "B": 1, + "D": 176, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.72, + "evolution_flashlib_ms": 2.050605, + "evolution_kernel_ms": 0.301759, + "evolution_speedup": 6.7955, + "evolution_tflops": 4.8926, + "label": "post_d895_d176_b1_n512_k8192_d176", + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "runtime_coverage": false, + "seed": 217604, + "source": "high_k_low_n" + }, + { + "B": 2, + "D": 176, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 3.0356, + "evolution_flashlib_ms": 0.486367, + "evolution_kernel_ms": 0.191839, + "evolution_speedup": 2.5353, + "evolution_tflops": 7.696, + "label": "post_d895_d176_b2_n2048_k1024_d176", + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "runtime_coverage": false, + "seed": 217602, + "source": "tailpad_boundary" + }, + { + "B": 4, + "D": 176, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 23.9979, + "evolution_flashlib_ms": 0.492175, + "evolution_kernel_ms": 0.195712, + "evolution_speedup": 2.5148, + "evolution_tflops": 60.3497, + "label": "post_d895_d176_b4_n8192_k1024_d176", + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "runtime_coverage": false, + "seed": 217603, + "source": "tailpad_boundary" + }, + { + "B": 1, + "D": 224, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0611, + "evolution_flashlib_ms": 0.480639, + "evolution_kernel_ms": 0.170496, + "evolution_speedup": 2.8191, + "evolution_tflops": 0.1722, + "label": "post_d895_d224_b1_n256_k256_d224", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 222404, + "source": "medium_small_boundary" + }, + { + "B": 1, + "D": 224, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.9081, + "evolution_flashlib_ms": 2.069245, + "evolution_kernel_ms": 0.236528, + "evolution_speedup": 8.7484, + "evolution_tflops": 7.9443, + "label": "post_d895_d224_b1_n512_k8192_d224", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 222403, + "source": "high_k_low_n" + }, + { + "B": 2, + "D": 224, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 3.9885, + "evolution_flashlib_ms": 0.47112, + "evolution_kernel_ms": 0.176704, + "evolution_speedup": 2.6662, + "evolution_tflops": 10.6339, + "label": "post_d895_d224_b2_n2048_k1024_d224", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 222401, + "source": "new_medium_gap" + }, + { + "B": 4, + "D": 224, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 27.1681, + "evolution_flashlib_ms": 0.553311, + "evolution_kernel_ms": 0.190783, + "evolution_speedup": 2.9002, + "evolution_tflops": 78.7929, + "label": "post_d895_d224_b4_n8192_k1024_d224", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 222402, + "source": "new_medium_gap" + }, + { + "B": 1, + "D": 288, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0649, + "evolution_flashlib_ms": 0.581567, + "evolution_kernel_ms": 0.171104, + "evolution_speedup": 3.3989, + "evolution_tflops": 0.2206, + "label": "post_d895_d288_b1_n256_k256_d288", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 228804, + "source": "medium_small_boundary" + }, + { + "B": 1, + "D": 288, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0421, + "evolution_flashlib_ms": 57.444236, + "evolution_kernel_ms": 0.273615, + "evolution_speedup": 209.9455, + "evolution_tflops": 8.8296, + "label": "post_d895_d288_b1_n512_k8192_d288", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 228803, + "source": "high_k_low_n" + }, + { + "B": 2, + "D": 288, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 3.7159, + "evolution_flashlib_ms": 0.65016, + "evolution_kernel_ms": 0.18616, + "evolution_speedup": 3.4925, + "evolution_tflops": 12.9777, + "label": "post_d895_d288_b2_n2048_k1024_d288", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 228801, + "source": "new_medium_gap" + }, + { + "B": 4, + "D": 288, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 28.8551, + "evolution_flashlib_ms": 0.669808, + "evolution_kernel_ms": 0.214687, + "evolution_speedup": 3.1199, + "evolution_tflops": 90.0257, + "label": "post_d895_d288_b4_n8192_k1024_d288", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 228802, + "source": "new_medium_gap" + }, + { + "B": 4, + "D": 32, + "K": 1024, + "N": 32768, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 23.8017, + "evolution_flashlib_ms": 0.360896, + "evolution_kernel_ms": 0.218816, + "evolution_speedup": 1.6493, + "evolution_tflops": 39.2564, + "label": "post_d895_d32_b4_n32768_k1024_d32", + "route": "microdim_hybrid_9c0d_v1", + "runtime_coverage": false, + "seed": 23202, + "source": "near_floor_microdim" + }, + { + "B": 8, + "D": 32, + "K": 512, + "N": 65536, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 34.3587, + "evolution_flashlib_ms": 0.500015, + "evolution_kernel_ms": 0.301439, + "evolution_speedup": 1.6588, + "evolution_tflops": 56.9928, + "label": "post_d895_d32_b8_n65536_k512_d32", + "route": "microdim_hybrid_9c0d_v1", + "runtime_coverage": false, + "seed": 23201, + "source": "near_floor_microdim" + }, + { + "B": 1, + "D": 352, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0736, + "evolution_flashlib_ms": 0.627167, + "evolution_kernel_ms": 0.172703, + "evolution_speedup": 3.6315, + "evolution_tflops": 0.2671, + "label": "post_d895_d352_b1_n256_k256_d352", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 235204, + "source": "high_small_boundary" + }, + { + "B": 1, + "D": 352, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0616, + "evolution_flashlib_ms": 47.910581, + "evolution_kernel_ms": 0.272767, + "evolution_speedup": 175.6465, + "evolution_tflops": 10.8253, + "label": "post_d895_d352_b1_n512_k8192_d352", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 235203, + "source": "high_k_low_n" + }, + { + "B": 2, + "D": 352, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 4.882, + "evolution_flashlib_ms": 0.604831, + "evolution_kernel_ms": 0.189184, + "evolution_speedup": 3.1971, + "evolution_tflops": 15.608, + "label": "post_d895_d352_b2_n2048_k1024_d352", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 235201, + "source": "new_high_gap" + }, + { + "B": 4, + "D": 352, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 36.2395, + "evolution_flashlib_ms": 0.651839, + "evolution_kernel_ms": 0.220992, + "evolution_speedup": 2.9496, + "evolution_tflops": 106.8922, + "label": "post_d895_d352_b4_n8192_k1024_d352", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 235202, + "source": "new_high_gap" + }, + { + "B": 1, + "D": 416, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0842, + "evolution_flashlib_ms": 0.647807, + "evolution_kernel_ms": 0.175137, + "evolution_speedup": 3.6989, + "evolution_tflops": 0.3113, + "label": "post_d895_d416_b1_n256_k256_d416", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 241604, + "source": "high_small_boundary" + }, + { + "B": 1, + "D": 416, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0852, + "evolution_flashlib_ms": 40.967689, + "evolution_kernel_ms": 0.275103, + "evolution_speedup": 148.9176, + "evolution_tflops": 12.6849, + "label": "post_d895_d416_b1_n512_k8192_d416", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 241603, + "source": "high_k_low_n" + }, + { + "B": 2, + "D": 416, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 5.4317, + "evolution_flashlib_ms": 0.642463, + "evolution_kernel_ms": 0.193983, + "evolution_speedup": 3.312, + "evolution_tflops": 17.9895, + "label": "post_d895_d416_b2_n2048_k1024_d416", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 241601, + "source": "new_high_gap" + }, + { + "B": 4, + "D": 416, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 39.3193, + "evolution_flashlib_ms": 0.710015, + "evolution_kernel_ms": 0.230463, + "evolution_speedup": 3.0808, + "evolution_tflops": 121.1354, + "label": "post_d895_d416_b4_n8192_k1024_d416", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 241602, + "source": "new_high_gap" + }, + { + "B": 1, + "D": 480, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0969, + "evolution_flashlib_ms": 0.649151, + "evolution_kernel_ms": 0.176655, + "evolution_speedup": 3.6747, + "evolution_tflops": 0.3561, + "label": "post_d895_d480_b1_n256_k256_d480", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 248004, + "source": "high_small_boundary" + }, + { + "B": 1, + "D": 480, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1226, + "evolution_flashlib_ms": 32.836689, + "evolution_kernel_ms": 0.278559, + "evolution_speedup": 117.8803, + "evolution_tflops": 14.4548, + "label": "post_d895_d480_b1_n512_k8192_d480", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 248003, + "source": "high_k_low_n" + }, + { + "B": 2, + "D": 480, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 5.9879, + "evolution_flashlib_ms": 0.672447, + "evolution_kernel_ms": 0.197919, + "evolution_speedup": 3.3976, + "evolution_tflops": 20.3443, + "label": "post_d895_d480_b2_n2048_k1024_d480", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 248001, + "source": "new_high_gap" + }, + { + "B": 4, + "D": 480, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 43.8869, + "evolution_flashlib_ms": 0.733983, + "evolution_kernel_ms": 0.24, + "evolution_speedup": 3.0583, + "evolution_tflops": 134.2177, + "label": "post_d895_d480_b4_n8192_k1024_d480", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 248002, + "source": "new_high_gap" + }, + { + "B": 1, + "D": 48, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.275, + "evolution_flashlib_ms": 1.464301, + "evolution_kernel_ms": 0.21952, + "evolution_speedup": 6.6705, + "evolution_tflops": 1.8342, + "label": "post_d895_d48_b1_n512_k8192_d48", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 24803, + "source": "high_k_low_n" + }, + { + "B": 2, + "D": 48, + "K": 512, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.6574, + "evolution_flashlib_ms": 0.30624, + "evolution_kernel_ms": 0.166656, + "evolution_speedup": 1.8376, + "evolution_tflops": 1.208, + "label": "post_d895_d48_b2_n2048_k512_d48", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 24801, + "source": "new_lowmid_gap" + }, + { + "B": 4, + "D": 48, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 9.5879, + "evolution_flashlib_ms": 0.335967, + "evolution_kernel_ms": 0.179807, + "evolution_speedup": 1.8685, + "evolution_tflops": 17.9149, + "label": "post_d895_d48_b4_n8192_k1024_d48", + "route": "gap_pad_to_supported_seed_v1", + "runtime_coverage": false, + "seed": 24802, + "source": "new_lowmid_gap" + } + ], + "submission_metric": { + "field": "submission_ms", + "host_enqueue_alias_field": "host_enqueue_ms", + "relationship": "same_cpu_submission_bracket" + }, + "timing_window_ms": { + "bench_ms": 100.0, + "warmup_ms": 20.0 + }, + "validation_shard": { + "count": 4, + "index": 0 + }, + "validation_shards": [ + { + "count": 4, + "index": 0 + }, + { + "count": 4, + "index": 1 + }, + { + "count": 4, + "index": 2 + }, + { + "count": 4, + "index": 3 + } + ] +} diff --git a/cake_exports/kmeans/EXPORT_TIMING.json b/cake_exports/kmeans/EXPORT_TIMING.json new file mode 100644 index 00000000..adab9ef1 --- /dev/null +++ b/cake_exports/kmeans/EXPORT_TIMING.json @@ -0,0 +1,11 @@ +{ + "host_wall_phase_seconds": { + "benchmark": 204.77159006567672, + "correctness": 18.721056894399226, + "validation_total": 223.49264696007594 + }, + "recorded_at_utc": "2026-07-09T04:30:10+00:00", + "reportable_gpu_timing_source": "BENCHMARK_RESULTS.json:CUPTI", + "schema": "loom-export-host-wall-phases-v1", + "validation_shard_count": 4 +} diff --git a/cake_exports/kmeans/README.md b/cake_exports/kmeans/README.md new file mode 100644 index 00000000..d879d320 --- /dev/null +++ b/cake_exports/kmeans/README.md @@ -0,0 +1,567 @@ +# Exported Loom Kernels + +## Pre-publication GPU validation: PASS — declared 124-shape performance floor + +- Hardware: `NVIDIA GB200` (`sm_100a`) +- Shapes: correctness `124/124`, CUPTI benchmark `124/124`; full generated suite `128` tests +- Validation shards: `4`; host wall time: correctness `18.72s`, benchmark `204.77s` +- `public_raw_e2e_speedup_vs_07cf_adapter` vs `triton_h200_07cf_raw_adapter_v1`: min `1.0938x`, geomean `2.1825x`, median `2.1942x`, p90 `2.8423x`, max `3.8773x` across all `124` floor-gated shapes (required minimum `1.0000x`) +- Candidate lifecycle latency diagnostics: init-once median `242.7617 ms`; first-signature compute median/p90 `3.1304/70.2117 ms`; hot compute median/p90 `0.0889/0.1581 ms` + +#### Hot steady-state synchronized E2E speedup + +| Validated shape scope | Min | Geomean | Median | P90 | Max | +| --- | ---: | ---: | ---: | ---: | ---: | +| All 124 benchmarked shapes (diagnostic scope) | 1.0938x | 2.1825x | 2.1942x | 2.8423x | 3.8773x | + +#### Modeled after-init amortized synchronized E2E speedup + +| Public calls N | Min | Geomean | Median | P90 | Max | +| --- | ---: | ---: | ---: | ---: | ---: | +| 1 | 0.0865x | 1.2309x | 2.7840x | 3.3516x | 163.7986x | +| 10 | 0.1110x | 1.2589x | 2.6667x | 3.0908x | 137.3500x | +| 100 | 0.2850x | 1.5093x | 2.1682x | 2.9894x | 53.5582x | +| 1000 | 1.0120x | 1.9957x | 2.0271x | 2.8424x | 9.5268x | + +#### Modeled including-init amortized synchronized E2E speedup + +| Public calls N | Min | Geomean | Median | P90 | Max | +| --- | ---: | ---: | ---: | ---: | ---: | +| 1 | 0.0772x | 0.9450x | 2.2810x | 12.9363x | 24.8909x | +| 10 | 0.0814x | 0.9591x | 2.2786x | 12.7562x | 24.5273x | +| 100 | 0.1093x | 1.0709x | 2.2565x | 10.9253x | 21.4483x | +| 1000 | 0.3510x | 1.4680x | 1.5030x | 6.0454x | 10.3790x | + +- All three tables report synchronized host E2E speedups as `baseline/candidate`. `Hot steady-state` measures a repeated public call at each lane's declared hot cache state; its per-shape values supply the official metric used by the separate publication-floor section. +- `After-init amortized(N) = (first_compute + (N-1) * hot_median) / N`; it excludes init. +- `Including-init amortized(N) = (init + first_compute + (N-1) * hot_median) / N`; it includes init. Each latency formula is evaluated separately for baseline and candidate, then reported as `baseline/candidate`. Both amortized scenarios are composed from measured components, not a directly timed N-call loop. A lane without explicit init uses `I=0`. +- Init scope: `runtime_init_plus_standalone_shared_preprocess_support_once_per_validation_shard_process_device_operator`; composition: `runtime_init_plus_shared_preprocess_support_each_standalone_lane`; baseline has explicit init: `yes`. +- Cache policy: `synchronize_and_clear_after_each_completed_shape`; resident multi-shape cache benchmarked: `no`; cold order: `deterministic_balanced_per_publication_contract_portfolio`; init order: `alternate_candidate_baseline_first_by_validation_shard_parity_then_shared_support` +- Lifecycle timing convention: all three lifecycle tables are synchronized host E2E. Init/first-call brackets are CUPTI timestamp host diagnostics; separately, the hot GPU-span diagnostic remains strict correlated CUPTI activity timing. + +- Measured: `2026-07-09T04:30:10+00:00` +- Full summary: [`VALIDATION.json`](VALIDATION.json); per-shape results: [`BENCHMARK_RESULTS.json`](BENCHMARK_RESULTS.json) + +This repository was generated by `loom.export.kernel_repo`. It contains frozen +CUDA source plus a lightweight Python binding that compiles the source with +NVRTC and launches it through CUDA driver APIs. + +## Complete Workload Entry Points + +This export is a complete workload package. Its three review and execution +entry points are deliberately separate and have stable names: + +- Python interface: [`src/flashlib_cake_kmeans/interface.py`](src/flashlib_cake_kmeans/interface.py) — public APIs: `FlashKMeansAssignRuntime`, `PreparedFlashKMeansAssign`, `init`, `prepare_flash_kmeans_assign`, `flash_kmeans_assign_prepared`, `flash_kmeans_assign` +- GPU correctness test: [`tests/test_correctness.py`](tests/test_correctness.py) +- CUPTI performance benchmark: [`benchmarks/benchmark.py`](benchmarks/benchmark.py) + +```bash +pytest tests/test_correctness.py -q +python benchmarks/benchmark.py --no-correctness --json results/performance.json +``` + +The public interface can be imported directly from the package root; the +implementation contains frozen CUDA and bindings but no Weave IR. + +## Repository Layout + +```text +src/flashlib_cake_kmeans/ + __init__.py # public Python exports + kernels.py # KernelSpec, get_kernel(), launch_() wrappers + tvm_ffi.py # optional TVM FFI global-function registration + _runtime.py # NVRTC + CUDA driver launch support + _benchmark.py # strict CUPTI timing with cold-L2 flushing + manifest.json # provenance, parameter order, launch metadata + cuda/*.cu # frozen bundled CUDA sources; no Weave IR +tests/ + test_exported_kernels.py + test_benchmark_harness.py +benchmarks/ + benchmark_exported_kernels.py + benchmark_shapes.py # semantic correctness + CUPTI performance runner + workload.py # workload shapes/reference/metric adapter +RESULTS.md # correctness and performance results +``` + +## Kernels + +- `dispatch_kernel_0000`: `kernel_flash_kmeans_assign_lowdim_pack_e50c_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0001`: `kernel_flash_kmeans_assign_lowdim_e50c_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0002`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0003`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0004`: `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0005`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0006`: `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0007`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0008`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0009`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0010`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0011`: `kernel_flash_kmeans_assign_highd_splitk_partial_blockn64_g2r4_b5a6_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_scores, partial_indices, B, N, D, K, num_n_tiles, K_tiles, K_slices) +- `dispatch_kernel_0012`: `kernel_flash_kmeans_assign_highd_splitk_reduce_blockn64_g2r4_b5a6_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_scores, partial_indices, out, B, N, K, num_n_tiles, K_slices) +- `dispatch_kernel_0013`: `kernel_flash_kmeans_assign_microdim_pack_6cd2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0014`: `kernel_flash_kmeans_assign_microdim_6cd2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0015`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0016`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0017`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0018`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0019`: `kernel_flash_kmeans_assign_microdim_pack_6cd2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0020`: `kernel_flash_kmeans_assign_microdim_6cd2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0021`: `kernel_flash_kmeans_assign_microdim_pack_6cd2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0022`: `kernel_flash_kmeans_assign_microdim_6cd2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0023`: `kernel_flash_kmeans_assign_microdim_pack_6cd2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0024`: `kernel_flash_kmeans_assign_microdim_6cd2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0025`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0026`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0027`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0028`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0029`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0030`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0031`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0032`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0033`: `kernel_flash_kmeans_assign_lowdim_pack_e50c_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0034`: `kernel_flash_kmeans_assign_lowdim_e50c_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0035`: `kernel_flash_kmeans_assign_lowdim_pack_e50c_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0036`: `kernel_flash_kmeans_assign_lowdim_e50c_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0037`: `kernel_flash_kmeans_assign_lowdim_pack_e50c_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0038`: `kernel_flash_kmeans_assign_lowdim_e50c_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0039`: `kernel_flash_kmeans_assign_lowdim_pack_e50c_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0040`: `kernel_flash_kmeans_assign_lowdim_e50c_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0041`: `kernel_flash_kmeans_assign_lowdim_pack_e50c_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0042`: `kernel_flash_kmeans_assign_lowdim_e50c_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0043`: `kernel_flash_kmeans_assign_lowdim_pack_e50c_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0044`: `kernel_flash_kmeans_assign_lowdim_e50c_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0045`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0046`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0047`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0048`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0049`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0050`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0051`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0052`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0053`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0054`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0055`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0056`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0057`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0058`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0059`: `kernel_flash_kmeans_assign_highd_splitk_partial_blockn64_g1r4_streamdep_r63_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_scores, partial_indices, B, N, D, K, num_n_tiles, K_tiles, K_slices) +- `dispatch_kernel_0060`: `kernel_flash_kmeans_assign_highd_splitk_reduce_blockn64_g1r4_streamdep_r63_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_scores, partial_indices, out, B, N, K, num_n_tiles, K_slices) +- `dispatch_kernel_0061`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0062`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0063`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0064`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0065`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0066`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0067`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0068`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0069`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0070`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0071`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0072`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0073`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0074`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0075`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0076`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0077`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0078`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0079`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0080`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0081`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0082`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0083`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0084`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0085`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0086`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0087`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0088`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0089`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0090`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0091`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0092`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0093`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0094`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0095`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0096`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0097`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0098`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0099`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0100`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0101`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0102`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0103`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0104`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0105`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0106`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0107`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0108`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0109`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0110`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0111`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0112`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0113`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0114`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0115`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0116`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0117`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0118`: `kernel_flash_kmeans_assign_highd_paired_xreuse_dualtmem_producer_r47_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_keys, B, N, D, K, num_n_tiles, K_tiles, K_slices) +- `dispatch_kernel_0119`: `kernel_flash_kmeans_assign_highd_paired_ownerreduce_r39_reduce1_unroll_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_keys, out, B, N, K, num_n_tiles, K_slices) +- `dispatch_kernel_0120`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0121`: `kernel_flash_kmeans_assign_highd_paired_packedpartial_producer_7b3c_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_keys, B, N, D, K, num_n_tiles, K_tiles, K_slices) +- `dispatch_kernel_0122`: `kernel_flash_kmeans_assign_highd_paired_packedpartial_reduce_r2_7b3c_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_keys, out, B, N, K, num_n_tiles, K_slices) +- `dispatch_kernel_0123`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0124`: `kernel_flash_kmeans_assign_microdim_pack_6cd2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0125`: `kernel_flash_kmeans_assign_microdim_6cd2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0126`: `kernel_flash_kmeans_assign_microdim_d16_pipeline4_08f9_v4`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0127`: `kernel_flash_kmeans_assign_microdim_pack_6cd2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0128`: `kernel_flash_kmeans_assign_microdim_6cd2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0129`: `kernel_flash_kmeans_assign_microdim_raw_tma_08f9_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0130`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0131`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0132`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0133`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0134`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0135`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0136`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0137`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0138`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0139`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0140`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0141`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0142`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0143`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0144`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0145`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0146`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0147`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0148`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0149`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0150`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0151`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0152`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0153`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0154`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0155`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0156`: `kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, c_sq, out, B, N, D, K, num_n_tiles) +- `dispatch_kernel_0157`: `kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, c_sq, out, B, N, D, K, num_n_tiles) +- `dispatch_kernel_0158`: `kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, c_sq, out, B, N, D, K, num_n_tiles) +- `dispatch_kernel_0159`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0160`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0161`: `kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, c_sq, out, B, N, D, K, num_n_tiles) +- `dispatch_kernel_0162`: `kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, c_sq, out, B, N, D, K, num_n_tiles) +- `dispatch_kernel_0163`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0164`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0165`: `kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, c_sq, out, B, N, D, K, num_n_tiles) +- `dispatch_kernel_0166`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0167`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0168`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0169`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0170`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0171`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0172`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0173`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0174`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0175`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0176`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0177`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0178`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0179`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0180`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0181`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0182`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0183`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0184`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0185`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0186`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0187`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0188`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0189`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0190`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0191`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0192`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0193`: `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0194`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0195`: `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0196`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0197`: `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0198`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0199`: `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0200`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0201`: `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0202`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0203`: `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0204`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0205`: `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0206`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0207`: `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0208`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0209`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0210`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0211`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0212`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0213`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0214`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0215`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0216`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0217`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0218`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0219`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0220`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0221`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0222`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0223`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0224`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0225`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0226`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0227`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0228`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0229`: `kernel_flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0230`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0231`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0232`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0233`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0234`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0235`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0236`: `kernel_flash_kmeans_assign_d288_exactd_a532_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0237`: `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_scores, partial_indices, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0238`: `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_scores, partial_indices, out, B, N, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0239`: `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_scores, partial_indices, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0240`: `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_scores, partial_indices, out, B, N, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0241`: `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_scores, partial_indices, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0242`: `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_scores, partial_indices, out, B, N, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0243`: `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_scores, partial_indices, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0244`: `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_scores, partial_indices, out, B, N, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0245`: `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_scores, partial_indices, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0246`: `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_scores, partial_indices, out, B, N, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0247`: `kernel_flash_kmeans_assign_d288_exactd_a532_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0248`: `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_scores, partial_indices, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0249`: `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_scores, partial_indices, out, B, N, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0250`: `kernel_flash_kmeans_assign_d288_exactd_a532_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0251`: `kernel_flash_kmeans_assign_d288_exactd_a532_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0252`: `kernel_flash_kmeans_assign_d288_exactd_a532_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0253`: `kernel_flash_kmeans_assign_d288_exactd_a532_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0254`: `kernel_flash_kmeans_assign_d288_exactd_a532_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0255`: `kernel_flash_kmeans_assign_d288_exactd_a532_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0256`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0257`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0258`: `kernel_flash_kmeans_assign_d480_splitk_partial_d32k256_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_scores, partial_indices, B, N, D, K, num_n_tiles, K_tiles, K_slices) +- `dispatch_kernel_0259`: `kernel_flash_kmeans_assign_d480_splitk_reduce_d32k256_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_scores, partial_indices, out, B, N, K, num_n_tiles, K_slices) +- `dispatch_kernel_0260`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0261`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0262`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0263`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0264`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0265`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0266`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0267`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0268`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0269`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0270`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0271`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0272`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0273`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0274`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0275`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0276`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0277`: `kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_scores, partial_indices, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0278`: `kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_scores, partial_indices, out, B, N, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0279`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0280`: `kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_scores, partial_indices, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0281`: `kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_scores, partial_indices, out, B, N, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0282`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0283`: `kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_scores, partial_indices, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0284`: `kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_scores, partial_indices, out, B, N, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0285`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0286`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0287`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0288`: `kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_scores, partial_indices, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0289`: `kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_scores, partial_indices, out, B, N, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0290`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0291`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0292`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0293`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0294`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0295`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0296`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0297`: `kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_scores, partial_indices, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0298`: `kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_scores, partial_indices, out, B, N, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0299`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0300`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0301`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0302`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0303`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0304`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0305`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0306`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0307`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0308`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0309`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0310`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0311`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0312`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0313`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0314`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0315`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0316`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0317`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0318`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0319`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0320`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0321`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0322`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0323`: `kernel_flash_kmeans_assign_highd_splitk_partial_blockn64_g2r4_b5a6_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_scores, partial_indices, B, N, D, K, num_n_tiles, K_tiles, K_slices) +- `dispatch_kernel_0324`: `kernel_flash_kmeans_assign_highd_splitk_reduce_blockn64_g2r4_b5a6_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_scores, partial_indices, out, B, N, K, num_n_tiles, K_slices) +- `dispatch_kernel_0325`: `kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_scores, partial_indices, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0326`: `kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_scores, partial_indices, out, B, N, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0327`: `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0328`: `kernel_flash_kmeans_assign_highd_splitk_partial_blockn64_g1r4_streamdep_r63_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_scores, partial_indices, B, N, D, K, num_n_tiles, K_tiles, K_slices) +- `dispatch_kernel_0329`: `kernel_flash_kmeans_assign_highd_splitk_reduce_blockn64_g1r4_streamdep_r63_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_scores, partial_indices, out, B, N, K, num_n_tiles, K_slices) +- `dispatch_kernel_0330`: `kernel_flash_kmeans_assign_microdim_d16_pipeline4_08f9_v4`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0331`: `kernel_flash_kmeans_assign_microdim_raw_tma_08f9_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0332`: `kernel_flash_kmeans_assign_microdim_direct_9c0d_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0333`: `kernel_flash_kmeans_assign_microdim_pack_6cd2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0334`: `kernel_flash_kmeans_assign_microdim_6cd2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0335`: `kernel_flash_kmeans_assign_lowdim_pack_e50c_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0336`: `kernel_flash_kmeans_assign_microdim_6cd2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0337`: `kernel_flash_kmeans_assign_lowdim_e50c_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0338`: `kernel_flash_kmeans_assign_gap_pad_pack_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, D_PAD, total_x_pad, total_c_pad) +- `dispatch_kernel_0339`: `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0340`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0341`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0342`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0343`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0344`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0345`: `kernel_flash_kmeans_assign_d480_splitk_partial_d32k256_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_scores, partial_indices, B, N, D, K, num_n_tiles, K_tiles, K_slices) +- `dispatch_kernel_0346`: `kernel_flash_kmeans_assign_d480_splitk_reduce_d32k256_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_scores, partial_indices, out, B, N, K, num_n_tiles, K_slices) +- `dispatch_kernel_0347`: `kernel_flash_kmeans_assign_d416_exactd_splitd_a4a579d1_v2`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0348`: `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, partial_scores, partial_indices, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0349`: `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce`, launch_mode=`standard`, threads=256, shared_mem=0, params=(partial_scores, partial_indices, out, B, N, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0350`: `kernel_flash_kmeans_assign_d288_exactd_a532_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0351`: `kernel_flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, x_sq, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0352`: `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) +- `dispatch_kernel_0353`: `kernel_flash_kmeans_assign_d112_k1024_owner_local_warp_mma_distinct_workfeed_c829_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, c_sq, out, B, N, D, K, num_n_tiles) +- `dispatch_kernel_0354`: `kernel_flash_kmeans_assign_d112_k512_two_owner_peer_only_mma_f826_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, c_sq, out, B, N, D, K, num_n_tiles) +- `dispatch_kernel_0355`: `kernel_flash_kmeans_assign_d112_shared_point_dual_issuer_tcgen05_6e1e_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x_tmap, c_tmap, c_sq, out, B, N, D, K, num_n_tiles, K_tiles) +- `dispatch_kernel_0356`: `kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, c_sq, out, B, N, D, K, num_n_tiles) +- `dispatch_kernel_0357`: `kernel_flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1`, launch_mode=`standard`, threads=256, shared_mem=0, params=(x, centroids, x_pad, c_pad, B, N, D, K, total_x_pad, total_c_pad) + +## Install + +```bash +python -m pip install -e . +``` + +The runtime expects PyTorch and a CUDA Python package that provides +`cuda.bindings.driver` and `cuda.bindings.nvrtc`. + +Install the optional Apache TVM FFI adapter with: + +```bash +python -m pip install -e ".[tvm-ffi]" +``` + +## API Entry Points + +The complete workload API is exported from the package root. The bindings below +remain available for inspecting or launching individual frozen kernels. + +- `flashlib_cake_kmeans.KERNELS`: mapping from exported kernel name to `ExportedKernel`. +- `flashlib_cake_kmeans.get_kernel(name)`: inspect metadata, source text, or compile. +- `flashlib_cake_kmeans.launch_(*args, grid=...)`: low-level launch wrapper. +- `ExportedKernel.source_text()`: read the frozen CUDA source. +- `ExportedKernel.compile(arch="sm_100a")`: compile with NVRTC. +- `ExportedKernel.launch(...)`: launch with explicit CUDA parameter order. + +Frozen sources are compiled through a process-local content-addressed cache. +Aliases with identical CUDA source, architecture, options, active device, and +function symbol share one loaded kernel; translation units with different +symbols or active device contexts share the NVRTC cubin but retain separate +loaded function objects. Semantic APIs may use +the generated `_runtime.launch_context(...)` helper to capture one current or +explicit PyTorch CUDA stream for every stage in a multi-kernel call. + +```python +import torch +from flashlib_cake_kmeans import get_kernel, launch_dispatch_kernel_0000 + +kernel = get_kernel("dispatch_kernel_0000") +print(kernel.parameters) + +# Pass CUDA tensors or scalar values in the parameter order shown in +# src/flashlib_cake_kmeans/manifest.json. Pointer parameters accept CUDA tensors +# or integer device pointers. Scalar parameters are range-checked and packed +# with the exact signedness and width recorded by the Cake-STD-derived ABI. +launch_dispatch_kernel_0000(*args, grid=(1, 1, 1)) +torch.cuda.synchronize() +``` + +## TVM FFI + +`flashlib_cake_kmeans.register_tvm_ffi()` registers both the workload-level public API +declared by a complete export plan and every low-level frozen kernel in the +TVM FFI global registry. Tensor arguments use DLPack zero-copy conversion and +launches honor the current TVM FFI CUDA stream. + +```python +import tvm_ffi +import flashlib_cake_kmeans + +names = flashlib_cake_kmeans.register_tvm_ffi() +print(names) + +# Complete plans expose semantic functions as .. +semantic = tvm_ffi.get_global_func("flashlib_cake_kmeans.", allow_missing=True) + +# Low-level functions use .launch_. The final seven +# positional arguments are grid xyz, block xyz, and dynamic shared-memory bytes. +launch = tvm_ffi.get_global_func("flashlib_cake_kmeans.launch_dispatch_kernel_0000") +launch(*kernel_args, 1, 1, 1, 256, 1, 1, 0) +``` + +## Correctness And Performance + +Recorded correctness and performance results live in `RESULTS.md`. +`benchmarks/workload.py` is the workload adapter: it declares concrete shapes +and builds candidate/reference/compare callables for each shape. When the +export command receives `--benchmark-adapter path/to/workload.py`, the +generated repository is immediately runnable; otherwise the file is an +explicit template that must be completed before performance can be reported. + +## Tests + +The exported repository includes unit tests that do not require a GPU or CUDA +Python bindings: + +```bash +python -m pip install -e ".[test]" +pytest +``` + +## Benchmarks + +The generated benchmark script measures NVRTC compile latency for the frozen +CUDA sources. It is a build/runtime smoke benchmark, not a semantic kernel +throughput benchmark: + +```bash +python benchmarks/benchmark_exported_kernels.py --arch sm_100a --json results/compile_benchmark.json +python benchmarks/benchmark_shapes.py --json results/shape_benchmark.json +``` + +The shape runner checks correctness before timing and refuses to report +performance unless the requested timing backend is CUPTI. `gpu_span_ms` remains +the official kernel timing used for throughput and speedup. Each result also +records the exact per-iteration kernel sum, active interval union, uncovered +gap, activity count, host enqueue bracket, synchronized end-to-end bracket, and +the first semantic-call host diagnostic. Host brackets may include an internal +API synchronization. Cold-L2 flushing completes before the host bracket starts, +so synchronized E2E isolates semantic call start through candidate completion. +Neither host bracket may be substituted for GPU-only timing. + +Stateful production plans may also publish an init-once lifecycle record. It +separates the one-time runtime init, first-signature compute, repeated hot +``runtime.compute`` distribution, and modeled amortized latency. Init and first +compute remain CUPTI-timestamp host diagnostics; only the repeated-call GPU span +is strict correlated CUPTI activity. The report states the output/preprocessing +policy, per-shard init samples, cache hit/miss evidence, clear policy, and whether +resident multi-shape caching was actually measured. diff --git a/cake_exports/kmeans/RESULTS.md b/cake_exports/kmeans/RESULTS.md new file mode 100644 index 00000000..7e7e114b --- /dev/null +++ b/cake_exports/kmeans/RESULTS.md @@ -0,0 +1,426 @@ +# Correctness And Performance Results + +## Export Provenance + +- Package: `flashlib_cake_kmeans` +- Source repository: `ssh://git@gitlab-master.nvidia.com:12051/cake/cake.git` +- Source commit: `42070e96d0734cb580854baef60f17625ba33bb5` +- Generated at: `2026-07-09T04:26:21.680464+00:00` + +## Latest Recorded Results + +## Pre-publication GPU validation: PASS — declared 124-shape performance floor + +- Hardware: `NVIDIA GB200` (`sm_100a`) +- Shapes: correctness `124/124`, CUPTI benchmark `124/124`; full generated suite `128` tests +- Validation shards: `4`; host wall time: correctness `18.72s`, benchmark `204.77s` +- `public_raw_e2e_speedup_vs_07cf_adapter` vs `triton_h200_07cf_raw_adapter_v1`: min `1.0938x`, geomean `2.1825x`, median `2.1942x`, p90 `2.8423x`, max `3.8773x` across all `124` floor-gated shapes (required minimum `1.0000x`) +- Candidate lifecycle latency diagnostics: init-once median `242.7617 ms`; first-signature compute median/p90 `3.1304/70.2117 ms`; hot compute median/p90 `0.0889/0.1581 ms` + +#### Hot steady-state synchronized E2E speedup + +| Validated shape scope | Min | Geomean | Median | P90 | Max | +| --- | ---: | ---: | ---: | ---: | ---: | +| All 124 benchmarked shapes (diagnostic scope) | 1.0938x | 2.1825x | 2.1942x | 2.8423x | 3.8773x | + +#### Modeled after-init amortized synchronized E2E speedup + +| Public calls N | Min | Geomean | Median | P90 | Max | +| --- | ---: | ---: | ---: | ---: | ---: | +| 1 | 0.0865x | 1.2309x | 2.7840x | 3.3516x | 163.7986x | +| 10 | 0.1110x | 1.2589x | 2.6667x | 3.0908x | 137.3500x | +| 100 | 0.2850x | 1.5093x | 2.1682x | 2.9894x | 53.5582x | +| 1000 | 1.0120x | 1.9957x | 2.0271x | 2.8424x | 9.5268x | + +#### Modeled including-init amortized synchronized E2E speedup + +| Public calls N | Min | Geomean | Median | P90 | Max | +| --- | ---: | ---: | ---: | ---: | ---: | +| 1 | 0.0772x | 0.9450x | 2.2810x | 12.9363x | 24.8909x | +| 10 | 0.0814x | 0.9591x | 2.2786x | 12.7562x | 24.5273x | +| 100 | 0.1093x | 1.0709x | 2.2565x | 10.9253x | 21.4483x | +| 1000 | 0.3510x | 1.4680x | 1.5030x | 6.0454x | 10.3790x | + +- All three tables report synchronized host E2E speedups as `baseline/candidate`. `Hot steady-state` measures a repeated public call at each lane's declared hot cache state; its per-shape values supply the official metric used by the separate publication-floor section. +- `After-init amortized(N) = (first_compute + (N-1) * hot_median) / N`; it excludes init. +- `Including-init amortized(N) = (init + first_compute + (N-1) * hot_median) / N`; it includes init. Each latency formula is evaluated separately for baseline and candidate, then reported as `baseline/candidate`. Both amortized scenarios are composed from measured components, not a directly timed N-call loop. A lane without explicit init uses `I=0`. +- Init scope: `runtime_init_plus_standalone_shared_preprocess_support_once_per_validation_shard_process_device_operator`; composition: `runtime_init_plus_shared_preprocess_support_each_standalone_lane`; baseline has explicit init: `yes`. +- Cache policy: `synchronize_and_clear_after_each_completed_shape`; resident multi-shape cache benchmarked: `no`; cold order: `deterministic_balanced_per_publication_contract_portfolio`; init order: `alternate_candidate_baseline_first_by_validation_shard_parity_then_shared_support` +- Lifecycle timing convention: all three lifecycle tables are synchronized host E2E. Init/first-call brackets are CUPTI timestamp host diagnostics; separately, the hot GPU-span diagnostic remains strict correlated CUPTI activity timing. + +- Measured: `2026-07-09T04:30:10+00:00` +- Full summary: [`VALIDATION.json`](VALIDATION.json); per-shape results: [`BENCHMARK_RESULTS.json`](BENCHMARK_RESULTS.json) + +The machine-readable per-shape public lifecycle and CUPTI evidence is in [`BENCHMARK_RESULTS.json`](BENCHMARK_RESULTS.json). + +## Result Table + +| Test | Hardware | Command | Result | +| --- | --- | --- | --- | +| metadata unit tests | not required | `pytest tests/test_exported_kernels.py tests/test_benchmark_harness.py -q` | pending | +| NVRTC compile benchmark | CUDA host | `python benchmarks/benchmark_exported_kernels.py --arch sm_100a --json results/compile_benchmark.json` | pending | +| semantic correctness | target GPU | `pytest tests/test_correctness.py -q` | PASS | +| kernel performance | target GPU | `python benchmarks/benchmark.py --no-correctness` | PASS (declared publication floor) | + +## Kernel Inventory + +| Name | Symbol | Launch Mode | Threads | Shared Memory Bytes | +| --- | --- | --- | ---: | ---: | +| `dispatch_kernel_0000` | `kernel_flash_kmeans_assign_lowdim_pack_e50c_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0001` | `kernel_flash_kmeans_assign_lowdim_e50c_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0002` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0003` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0004` | `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0005` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0006` | `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0007` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0008` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0009` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0010` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0011` | `kernel_flash_kmeans_assign_highd_splitk_partial_blockn64_g2r4_b5a6_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0012` | `kernel_flash_kmeans_assign_highd_splitk_reduce_blockn64_g2r4_b5a6_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0013` | `kernel_flash_kmeans_assign_microdim_pack_6cd2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0014` | `kernel_flash_kmeans_assign_microdim_6cd2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0015` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0016` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0017` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0018` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0019` | `kernel_flash_kmeans_assign_microdim_pack_6cd2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0020` | `kernel_flash_kmeans_assign_microdim_6cd2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0021` | `kernel_flash_kmeans_assign_microdim_pack_6cd2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0022` | `kernel_flash_kmeans_assign_microdim_6cd2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0023` | `kernel_flash_kmeans_assign_microdim_pack_6cd2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0024` | `kernel_flash_kmeans_assign_microdim_6cd2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0025` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0026` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0027` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0028` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0029` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0030` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0031` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0032` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0033` | `kernel_flash_kmeans_assign_lowdim_pack_e50c_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0034` | `kernel_flash_kmeans_assign_lowdim_e50c_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0035` | `kernel_flash_kmeans_assign_lowdim_pack_e50c_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0036` | `kernel_flash_kmeans_assign_lowdim_e50c_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0037` | `kernel_flash_kmeans_assign_lowdim_pack_e50c_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0038` | `kernel_flash_kmeans_assign_lowdim_e50c_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0039` | `kernel_flash_kmeans_assign_lowdim_pack_e50c_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0040` | `kernel_flash_kmeans_assign_lowdim_e50c_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0041` | `kernel_flash_kmeans_assign_lowdim_pack_e50c_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0042` | `kernel_flash_kmeans_assign_lowdim_e50c_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0043` | `kernel_flash_kmeans_assign_lowdim_pack_e50c_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0044` | `kernel_flash_kmeans_assign_lowdim_e50c_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0045` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0046` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0047` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0048` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0049` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0050` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0051` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0052` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0053` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0054` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0055` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0056` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0057` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0058` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0059` | `kernel_flash_kmeans_assign_highd_splitk_partial_blockn64_g1r4_streamdep_r63_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0060` | `kernel_flash_kmeans_assign_highd_splitk_reduce_blockn64_g1r4_streamdep_r63_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0061` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0062` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0063` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0064` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0065` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0066` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0067` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0068` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0069` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0070` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0071` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0072` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0073` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0074` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0075` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0076` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0077` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0078` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0079` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0080` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0081` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0082` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0083` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0084` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0085` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0086` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0087` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0088` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0089` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0090` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0091` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0092` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0093` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0094` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0095` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0096` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0097` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0098` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0099` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0100` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0101` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0102` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0103` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0104` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0105` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0106` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0107` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0108` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0109` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0110` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0111` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0112` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0113` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0114` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0115` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0116` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0117` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0118` | `kernel_flash_kmeans_assign_highd_paired_xreuse_dualtmem_producer_r47_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0119` | `kernel_flash_kmeans_assign_highd_paired_ownerreduce_r39_reduce1_unroll_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0120` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0121` | `kernel_flash_kmeans_assign_highd_paired_packedpartial_producer_7b3c_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0122` | `kernel_flash_kmeans_assign_highd_paired_packedpartial_reduce_r2_7b3c_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0123` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0124` | `kernel_flash_kmeans_assign_microdim_pack_6cd2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0125` | `kernel_flash_kmeans_assign_microdim_6cd2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0126` | `kernel_flash_kmeans_assign_microdim_d16_pipeline4_08f9_v4` | `standard` | 256 | 0 | +| `dispatch_kernel_0127` | `kernel_flash_kmeans_assign_microdim_pack_6cd2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0128` | `kernel_flash_kmeans_assign_microdim_6cd2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0129` | `kernel_flash_kmeans_assign_microdim_raw_tma_08f9_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0130` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0131` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0132` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0133` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0134` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0135` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0136` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0137` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0138` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0139` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0140` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0141` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0142` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0143` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0144` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0145` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0146` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0147` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0148` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0149` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0150` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0151` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0152` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0153` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0154` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0155` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0156` | `kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0157` | `kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0158` | `kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0159` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0160` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0161` | `kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0162` | `kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0163` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0164` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0165` | `kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0166` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0167` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0168` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0169` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0170` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0171` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0172` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0173` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0174` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0175` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0176` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0177` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0178` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0179` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0180` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0181` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0182` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0183` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0184` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0185` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0186` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0187` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0188` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0189` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0190` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0191` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0192` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0193` | `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0194` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0195` | `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0196` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0197` | `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0198` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0199` | `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0200` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0201` | `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0202` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0203` | `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0204` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0205` | `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0206` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0207` | `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0208` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0209` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0210` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0211` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0212` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0213` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0214` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0215` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0216` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0217` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0218` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0219` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0220` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0221` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0222` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0223` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0224` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0225` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0226` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0227` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0228` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0229` | `kernel_flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4` | `standard` | 256 | 0 | +| `dispatch_kernel_0230` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0231` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0232` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0233` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0234` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0235` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0236` | `kernel_flash_kmeans_assign_d288_exactd_a532_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0237` | `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial` | `standard` | 256 | 0 | +| `dispatch_kernel_0238` | `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce` | `standard` | 256 | 0 | +| `dispatch_kernel_0239` | `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial` | `standard` | 256 | 0 | +| `dispatch_kernel_0240` | `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce` | `standard` | 256 | 0 | +| `dispatch_kernel_0241` | `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial` | `standard` | 256 | 0 | +| `dispatch_kernel_0242` | `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce` | `standard` | 256 | 0 | +| `dispatch_kernel_0243` | `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial` | `standard` | 256 | 0 | +| `dispatch_kernel_0244` | `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce` | `standard` | 256 | 0 | +| `dispatch_kernel_0245` | `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial` | `standard` | 256 | 0 | +| `dispatch_kernel_0246` | `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce` | `standard` | 256 | 0 | +| `dispatch_kernel_0247` | `kernel_flash_kmeans_assign_d288_exactd_a532_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0248` | `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial` | `standard` | 256 | 0 | +| `dispatch_kernel_0249` | `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce` | `standard` | 256 | 0 | +| `dispatch_kernel_0250` | `kernel_flash_kmeans_assign_d288_exactd_a532_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0251` | `kernel_flash_kmeans_assign_d288_exactd_a532_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0252` | `kernel_flash_kmeans_assign_d288_exactd_a532_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0253` | `kernel_flash_kmeans_assign_d288_exactd_a532_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0254` | `kernel_flash_kmeans_assign_d288_exactd_a532_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0255` | `kernel_flash_kmeans_assign_d288_exactd_a532_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0256` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0257` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0258` | `kernel_flash_kmeans_assign_d480_splitk_partial_d32k256_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0259` | `kernel_flash_kmeans_assign_d480_splitk_reduce_d32k256_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0260` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0261` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0262` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0263` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0264` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0265` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0266` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0267` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0268` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0269` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0270` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0271` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0272` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0273` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0274` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0275` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0276` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0277` | `kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0278` | `kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0279` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0280` | `kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0281` | `kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0282` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0283` | `kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0284` | `kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0285` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0286` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0287` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0288` | `kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0289` | `kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0290` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0291` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0292` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0293` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0294` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0295` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0296` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0297` | `kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0298` | `kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0299` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0300` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0301` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0302` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0303` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0304` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0305` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0306` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0307` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0308` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0309` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0310` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0311` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0312` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0313` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0314` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0315` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0316` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0317` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0318` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0319` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0320` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0321` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0322` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0323` | `kernel_flash_kmeans_assign_highd_splitk_partial_blockn64_g2r4_b5a6_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0324` | `kernel_flash_kmeans_assign_highd_splitk_reduce_blockn64_g2r4_b5a6_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0325` | `kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0326` | `kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0327` | `kernel_flash_kmeans_assign_highd_splitd_6fcf_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0328` | `kernel_flash_kmeans_assign_highd_splitk_partial_blockn64_g1r4_streamdep_r63_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0329` | `kernel_flash_kmeans_assign_highd_splitk_reduce_blockn64_g1r4_streamdep_r63_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0330` | `kernel_flash_kmeans_assign_microdim_d16_pipeline4_08f9_v4` | `standard` | 256 | 0 | +| `dispatch_kernel_0331` | `kernel_flash_kmeans_assign_microdim_raw_tma_08f9_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0332` | `kernel_flash_kmeans_assign_microdim_direct_9c0d_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0333` | `kernel_flash_kmeans_assign_microdim_pack_6cd2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0334` | `kernel_flash_kmeans_assign_microdim_6cd2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0335` | `kernel_flash_kmeans_assign_lowdim_pack_e50c_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0336` | `kernel_flash_kmeans_assign_microdim_6cd2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0337` | `kernel_flash_kmeans_assign_lowdim_e50c_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0338` | `kernel_flash_kmeans_assign_gap_pad_pack_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0339` | `kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0340` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v15` | `standard` | 256 | 0 | +| `dispatch_kernel_0341` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_v10` | `standard` | 256 | 0 | +| `dispatch_kernel_0342` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0343` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0344` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0345` | `kernel_flash_kmeans_assign_d480_splitk_partial_d32k256_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0346` | `kernel_flash_kmeans_assign_d480_splitk_reduce_d32k256_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0347` | `kernel_flash_kmeans_assign_d416_exactd_splitd_a4a579d1_v2` | `standard` | 256 | 0 | +| `dispatch_kernel_0348` | `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial` | `standard` | 256 | 0 | +| `dispatch_kernel_0349` | `kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce` | `standard` | 256 | 0 | +| `dispatch_kernel_0350` | `kernel_flash_kmeans_assign_d288_exactd_a532_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0351` | `kernel_flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4` | `standard` | 256 | 0 | +| `dispatch_kernel_0352` | `kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0353` | `kernel_flash_kmeans_assign_d112_k1024_owner_local_warp_mma_distinct_workfeed_c829_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0354` | `kernel_flash_kmeans_assign_d112_k512_two_owner_peer_only_mma_f826_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0355` | `kernel_flash_kmeans_assign_d112_shared_point_dual_issuer_tcgen05_6e1e_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0356` | `kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1` | `standard` | 256 | 0 | +| `dispatch_kernel_0357` | `kernel_flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1` | `standard` | 256 | 0 | diff --git a/cake_exports/kmeans/VALIDATION.json b/cake_exports/kmeans/VALIDATION.json new file mode 100644 index 00000000..e66206ff --- /dev/null +++ b/cake_exports/kmeans/VALIDATION.json @@ -0,0 +1,427 @@ +{ + "all_shape_diagnostics": { + "below_minimum_speedup": [], + "below_minimum_speedup_count": 0, + "floor_enforced": true, + "shape_count": 124, + "speedup_vs_baseline": { + "geomean": 2.1825046815076243, + "max": 3.8773006134969323, + "median": 2.1942486496321694, + "min": 1.093838998831605, + "p90": 2.842345086787142, + "sample_count": 124 + } + }, + "benchmark_commands": [ + [ + "/home/averyh/cake-export-venv-v6-kmeans/bin/python", + "benchmarks/benchmark.py", + "--no-correctness", + "--json", + "BENCHMARK_RESULTS.shard-0000-of-0004.json", + "--shard-index", + "0", + "--shard-count", + "4", + "--quiet" + ], + [ + "/home/averyh/cake-export-venv-v6-kmeans/bin/python", + "benchmarks/benchmark.py", + "--no-correctness", + "--json", + "BENCHMARK_RESULTS.shard-0001-of-0004.json", + "--shard-index", + "1", + "--shard-count", + "4", + "--quiet" + ], + [ + "/home/averyh/cake-export-venv-v6-kmeans/bin/python", + "benchmarks/benchmark.py", + "--no-correctness", + "--json", + "BENCHMARK_RESULTS.shard-0002-of-0004.json", + "--shard-index", + "2", + "--shard-count", + "4", + "--quiet" + ], + [ + "/home/averyh/cake-export-venv-v6-kmeans/bin/python", + "benchmarks/benchmark.py", + "--no-correctness", + "--json", + "BENCHMARK_RESULTS.shard-0003-of-0004.json", + "--shard-index", + "3", + "--shard-count", + "4", + "--quiet" + ] + ], + "benchmark_passed": 124, + "correctness_collection_command": [ + "/home/averyh/cake-export-venv-v6-kmeans/bin/python", + "-m", + "pytest", + "tests/test_correctness.py", + "--collect-only", + "-q", + "-m", + "export_validation_shape" + ], + "correctness_command": [ + "/home/averyh/cake-export-venv-v6-kmeans/bin/python", + "-m", + "pytest", + "tests/test_correctness.py", + "-q" + ], + "correctness_passed": 124, + "correctness_shape_count": 124, + "correctness_tests_passed": 128, + "hardware": { + "arch": "sm_100a", + "device": "NVIDIA GB200" + }, + "measured_at_utc": "2026-07-09T04:30:10+00:00", + "minimum_speedup": 1.0, + "performance_shape_count": 124, + "performance_speedup_denominator_metric": "candidate_public_raw_synchronized_e2e_ms", + "performance_speedup_metric": "public_raw_e2e_speedup_vs_07cf_adapter", + "performance_speedup_numerator_metric": "baseline_07cf_adapter_synchronized_e2e_ms", + "phase_seconds": { + "benchmark": 204.77159006567672, + "correctness": 18.721056894399226, + "validation_total": 223.49264696007594 + }, + "publication_performance_floor": { + "below_floor": [], + "below_floor_count": 0, + "minimum_speedup": 1.0, + "scope": "all_benchmarked_shapes", + "shape_count": 124, + "shape_labels": [ + "adjacent_1d49_d112_highk_tail_b1_n1408_k4096_d112", + "adjacent_1d49_d128_forced_fallback_b5_n7296_k512_d128", + "adjacent_1d49_d224_highn_overlap_b5_n5632_k768_d224", + "adjacent_1d49_d288_low_n_heldout_b3_n1152_k2048_d288", + "adjacent_1d49_d352_request_highk_b2_n1024_k8192_d352", + "adjacent_1d49_d416_random_b4_n3840_k512_d416", + "adjacent_1d49_d480_smallk_boundary_b2_n2816_k256_d480", + "adjacent_1d49_d48_lowk_boundary_b4_n3968_k256_d48", + "adjacent_3328_d112_highk_low_n_b2_n512_k8192_d112", + "adjacent_3328_d128_exact_guard_k_neighbor_b8_n8192_k512_d128", + "adjacent_3328_d224_tail_div_b3_n3840_k512_d224", + "adjacent_3328_d288_guard_overlap_b4_n8192_k256_d288", + "adjacent_3328_d352_random_legal_b3_n2048_k768_d352", + "adjacent_3328_d416_highk_low_n_b1_n384_k8192_d416", + "adjacent_3328_d480_min_boundary_b1_n128_k256_d480", + "adjacent_3328_d48_small_boundary_b1_n256_k256_d48", + "adjacent_5600_d112_odd_tile_tail_b5_n2176_k512_d112", + "adjacent_5600_d128_fallback_neighbor_b8_n8320_k256_d128", + "adjacent_5600_d224_k512_overlap_b3_n3072_k512_d224", + "adjacent_5600_d288_mid_highk_b2_n1024_k2048_d288", + "adjacent_5600_d352_splitk_edge_b2_n768_k4096_d352", + "adjacent_5600_d416_smallk_boundary_b4_n2048_k256_d416", + "adjacent_5600_d480_random_b1_n1536_k1024_d480", + "adjacent_5600_d48_low_n_boundary_b3_n1536_k256_d48", + "adjacent_68cf_d112_tail_b5_n2944_k512_d112", + "adjacent_68cf_d128_forced_fallback_b7_n6016_k512_d128", + "adjacent_68cf_d224_overlap_b3_n5120_k1024_d224", + "adjacent_68cf_d288_boundary_b2_n1920_k512_d288", + "adjacent_68cf_d352_tail_b3_n2816_k768_d352", + "adjacent_68cf_d416_overlap_b2_n2304_k1024_d416", + "adjacent_68cf_d480_boundary_b4_n1664_k512_d480", + "adjacent_68cf_d48_boundary_b4_n2304_k512_d48", + "adjacent_8f09_d112_small_highk_b1_n384_k4096_d112", + "adjacent_8f09_d128_fallback_neighbor_b4_n4480_k512_d128", + "adjacent_8f09_d224_low_n_boundary_b4_n1536_k256_d224", + "adjacent_8f09_d288_k768_overlap_b1_n2560_k768_d288", + "adjacent_8f09_d352_smallk_large_n_b4_n4096_k256_d352", + "adjacent_8f09_d416_midk_tail_b3_n3456_k768_d416", + "adjacent_8f09_d480_highk_low_n_b2_n640_k4096_d480", + "adjacent_8f09_d48_medium_tail_b2_n1792_k512_d48", + "adjacent_9ca0_d112_tail_div_b3_n3840_k768_d112", + "adjacent_9ca0_d128_exact_guard_neighbor_b8_n8064_k256_d128", + "adjacent_9ca0_d224_guard_overlap_b2_n4096_k256_d224", + "adjacent_9ca0_d288_highk_low_n_b1_n384_k4096_d288", + "adjacent_9ca0_d352_splitk_boundary_b1_n512_k8192_d352", + "adjacent_9ca0_d416_splitk_min_b1_n1024_k4096_d416", + "adjacent_9ca0_d480_splitd_large_n_b2_n4096_k512_d480", + "adjacent_9ca0_d48_guard_boundary_b6_n12288_k512_d48", + "adjacent_a2f8_d112_lowk_tail_b4_n3456_k256_d112", + "adjacent_a2f8_d128_odd_fallback_b6_n8576_k512_d128", + "adjacent_a2f8_d224_highn_overlap_b2_n6144_k768_d224", + "adjacent_a2f8_d288_splitk_probe_b1_n640_k4096_d288", + "adjacent_a2f8_d352_smallk_boundary_b4_n1024_k256_d352", + "adjacent_a2f8_d416_random_b2_n2560_k768_d416", + "adjacent_a2f8_d480_highk_request_b1_n896_k4096_d480", + "adjacent_a2f8_d48_k1024_boundary_b2_n768_k1024_d48", + "adjacent_c44f_d112_odd_tail_b2_n3200_k768_d112", + "adjacent_c44f_d128_forced_fallback_b6_n6272_k512_d128", + "adjacent_c44f_d224_overlap_b4_n4480_k512_d224", + "adjacent_c44f_d288_highk_heldout_b2_n768_k4096_d288", + "adjacent_c44f_d352_random_b1_n3328_k768_d352", + "adjacent_c44f_d416_highk_request_b3_n512_k8192_d416", + "adjacent_c44f_d480_boundary_b5_n2048_k512_d480", + "adjacent_c44f_d48_midk_boundary_b3_n2688_k768_d48", + "adjacent_d9d5_d112_highk_request_b3_n768_k8192_d112", + "adjacent_d9d5_d128_guard_miss_fallback_b5_n6016_k1024_d128", + "adjacent_d9d5_d224_random_midk_b2_n2944_k1280_d224", + "adjacent_d9d5_d288_heldout_low_n_b1_n896_k4096_d288", + "adjacent_d9d5_d352_random_b5_n2304_k768_d352", + "adjacent_d9d5_d416_highk_request_b2_n640_k8192_d416", + "adjacent_d9d5_d480_heldout_highk_b3_n1024_k4096_d480", + "adjacent_d9d5_d48_highk_heldout_b2_n640_k4096_d48", + "adjacent_ef00_d112_odd_tail_b4_n3712_k1024_d112", + "adjacent_ef00_d128_forced_fallback_b4_n7552_k512_d128", + "adjacent_ef00_d224_midn_overlap_b2_n5376_k512_d224", + "adjacent_ef00_d288_highk_heldout_b1_n1664_k4096_d288", + "adjacent_ef00_d352_request_low_n_b3_n896_k8192_d352", + "adjacent_ef00_d416_random_midk_b1_n2176_k768_d416", + "adjacent_ef00_d480_lowk_boundary_b3_n3200_k256_d480", + "adjacent_ef00_d48_midn_boundary_b3_n3456_k512_d48", + "post_d895_d112_b1_n256_k256_d112", + "post_d895_d112_b1_n512_k8192_d112", + "post_d895_d112_b2_n2048_k512_d112", + "post_d895_d112_b4_n8192_k1024_d112", + "post_d895_d128_fallback_b3_n1920_k256_d128", + "post_d895_d128_fallback_b5_n2176_k512_d128", + "post_d895_d128_fallback_b7_n2432_k1024_d128", + "post_d895_d128_paired_b2_n262144_k256_d128", + "post_d895_d128_paired_b8_n8192_k256_d128", + "post_d895_d144_b1_n256_k256_d144", + "post_d895_d144_b1_n512_k8192_d144", + "post_d895_d144_b2_n2048_k1024_d144", + "post_d895_d144_b4_n8192_k1024_d144", + "post_d895_d16_b4_n32768_k1024_d16", + "post_d895_d16_b8_n65536_k512_d16", + "post_d895_d176_b1_n256_k256_d176", + "post_d895_d176_b1_n512_k8192_d176", + "post_d895_d176_b2_n2048_k1024_d176", + "post_d895_d176_b4_n8192_k1024_d176", + "post_d895_d224_b1_n256_k256_d224", + "post_d895_d224_b1_n512_k8192_d224", + "post_d895_d224_b2_n2048_k1024_d224", + "post_d895_d224_b4_n8192_k1024_d224", + "post_d895_d288_b1_n256_k256_d288", + "post_d895_d288_b1_n512_k8192_d288", + "post_d895_d288_b2_n2048_k1024_d288", + "post_d895_d288_b4_n8192_k1024_d288", + "post_d895_d32_b4_n32768_k1024_d32", + "post_d895_d32_b8_n65536_k512_d32", + "post_d895_d352_b1_n256_k256_d352", + "post_d895_d352_b1_n512_k8192_d352", + "post_d895_d352_b2_n2048_k1024_d352", + "post_d895_d352_b4_n8192_k1024_d352", + "post_d895_d416_b1_n256_k256_d416", + "post_d895_d416_b1_n512_k8192_d416", + "post_d895_d416_b2_n2048_k1024_d416", + "post_d895_d416_b4_n8192_k1024_d416", + "post_d895_d480_b1_n256_k256_d480", + "post_d895_d480_b1_n512_k8192_d480", + "post_d895_d480_b2_n2048_k1024_d480", + "post_d895_d480_b4_n8192_k1024_d480", + "post_d895_d48_b1_n512_k8192_d48", + "post_d895_d48_b2_n2048_k512_d48", + "post_d895_d48_b4_n8192_k1024_d48" + ], + "speedup_vs_baseline": { + "geomean": 2.1825046815076243, + "max": 3.8773006134969323, + "median": 2.1942486496321694, + "min": 1.093838998831605, + "p90": 2.842345086787142, + "sample_count": 124 + }, + "status": "passed" + }, + "route_count": 124, + "runtime_lifecycle": { + "amortized_after_init_synchronized_e2e_speedup": { + "1": { + "geomean": 1.2309434323137933, + "max": 163.79859302023343, + "median": 2.7839934141126323, + "min": 0.08647202905817367, + "p90": 3.351566582003962, + "sample_count": 124 + }, + "10": { + "geomean": 1.2588822989268478, + "max": 137.34998767211528, + "median": 2.6667329239063666, + "min": 0.11096680845482784, + "p90": 3.090755401506125, + "sample_count": 124 + }, + "100": { + "geomean": 1.509344986064645, + "max": 53.5581614291322, + "median": 2.1681602849557455, + "min": 0.28495469134404755, + "p90": 2.989410840698895, + "sample_count": 124 + }, + "1000": { + "geomean": 1.9957356325578275, + "max": 9.526755432882533, + "median": 2.02714552137833, + "min": 1.0119821488655745, + "p90": 2.8424059645989446, + "sample_count": 124 + } + }, + "amortized_including_init_synchronized_e2e_speedup": { + "1": { + "geomean": 0.9449946730011419, + "max": 24.89092858456689, + "median": 2.2809542656641058, + "min": 0.07723037119783416, + "p90": 12.93634198448204, + "sample_count": 124 + }, + "10": { + "geomean": 0.9590724810364581, + "max": 24.52727808638291, + "median": 2.2785747135963885, + "min": 0.08144088410284854, + "p90": 12.756198569243546, + "sample_count": 124 + }, + "100": { + "geomean": 1.0709455274785604, + "max": 21.448324006518714, + "median": 2.256511456266482, + "min": 0.1092928565392167, + "p90": 10.92530869455313, + "sample_count": 124 + }, + "1000": { + "geomean": 1.4680428292344323, + "max": 10.378991521881256, + "median": 1.5029920226671636, + "min": 0.3509947998170528, + "p90": 6.045445059831923, + "sample_count": 124 + } + }, + "baseline_has_explicit_init": true, + "baseline_hot_public_synchronized_e2e_ms": { + "geomean": 0.20757016572840578, + "max": 0.5228809999999999, + "median": 0.18443199999999998, + "min": 0.16128, + "p90": 0.32694080000000014, + "sample_count": 124 + }, + "baseline_init_host_enqueue_ms": { + "geomean": 35.27461542656189, + "max": 36.934438, + "median": 35.000356, + "min": 34.227907, + "p90": 36.519487, + "sample_count": 4 + }, + "baseline_init_synchronized_e2e_ms": { + "geomean": 128.13715080146187, + "max": 484.041494, + "median": 245.37127800000002, + "min": 34.345443, + "p90": 475.36244150000005, + "sample_count": 4 + }, + "cache_policy": "synchronize_and_clear_after_each_completed_shape", + "candidate_first_compute_host_enqueue_ms": { + "geomean": 7.402417492079541, + "max": 83.722038, + "median": 3.0477155, + "min": 1.974882, + "p90": 70.15640230000001, + "sample_count": 124 + }, + "candidate_first_compute_synchronized_e2e_ms": { + "geomean": 7.475266319780226, + "max": 83.779798, + "median": 3.1303549999999998, + "min": 2.05613, + "p90": 70.21168870000001, + "sample_count": 124 + }, + "candidate_hot_compute_gpu_span_ms": { + "geomean": 0.03929959622925097, + "max": 0.389696, + "median": 0.038976, + "min": 0.009792, + "p90": 0.1054144, + "sample_count": 124 + }, + "candidate_hot_compute_host_enqueue_ms": { + "geomean": 0.04880204368462327, + "max": 0.06767999999999999, + "median": 0.047231999999999996, + "min": 0.040064, + "p90": 0.0595648, + "sample_count": 124 + }, + "candidate_hot_compute_synchronized_e2e_ms": { + "geomean": 0.09510640113954803, + "max": 0.438208, + "median": 0.08894425, + "min": 0.056224, + "p90": 0.1581216, + "sample_count": 124 + }, + "candidate_init_host_enqueue_ms": { + "geomean": 35.64562210040985, + "max": 37.358022999999996, + "median": 35.356052, + "min": 34.57754, + "p90": 36.902022099999996, + "sample_count": 4 + }, + "candidate_init_synchronized_e2e_ms": { + "geomean": 129.7884798704558, + "max": 484.100213, + "median": 242.7617075, + "min": 34.943076, + "p90": 473.29648670000006, + "sample_count": 4 + }, + "cold_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "hot_gpu_span_speedup": { + "geomean": 2.316300717543746, + "max": 6.406190408085431, + "median": 2.159249130548963, + "min": 0.9709337534898999, + "p90": 4.868219371609404, + "sample_count": 124 + }, + "hot_synchronized_e2e_speedup": { + "geomean": 2.1825046815076243, + "max": 3.8773006134969323, + "median": 2.1942486496321694, + "min": 1.093838998831605, + "p90": 2.842345086787142, + "sample_count": 124 + }, + "init_composition": "runtime_init_plus_shared_preprocess_support_each_standalone_lane", + "init_order_policy": "alternate_candidate_baseline_first_by_validation_shard_parity_then_shared_support", + "init_scope": "runtime_init_plus_standalone_shared_preprocess_support_once_per_validation_shard_process_device_operator", + "resident_multi_shape_cache_benchmarked": false, + "schema": "loom-public-runtime-lifecycle-summary-v1", + "shape_count": 124, + "validation_shard_count": 4 + }, + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape_count": 124, + "status": "passed", + "status_scope": "full_correctness_and_evidence_plus_declared_publication_performance_floor", + "timing_backend": "cupti", + "validation_shard_count": 4 +} diff --git a/cake_exports/kmeans/VALIDATION.md b/cake_exports/kmeans/VALIDATION.md new file mode 100644 index 00000000..81d24ced --- /dev/null +++ b/cake_exports/kmeans/VALIDATION.md @@ -0,0 +1,41 @@ +## Pre-publication GPU validation: PASS — declared 124-shape performance floor + +- Hardware: `NVIDIA GB200` (`sm_100a`) +- Shapes: correctness `124/124`, CUPTI benchmark `124/124`; full generated suite `128` tests +- Validation shards: `4`; host wall time: correctness `18.72s`, benchmark `204.77s` +- `public_raw_e2e_speedup_vs_07cf_adapter` vs `triton_h200_07cf_raw_adapter_v1`: min `1.0938x`, geomean `2.1825x`, median `2.1942x`, p90 `2.8423x`, max `3.8773x` across all `124` floor-gated shapes (required minimum `1.0000x`) +- Candidate lifecycle latency diagnostics: init-once median `242.7617 ms`; first-signature compute median/p90 `3.1304/70.2117 ms`; hot compute median/p90 `0.0889/0.1581 ms` + +#### Hot steady-state synchronized E2E speedup + +| Validated shape scope | Min | Geomean | Median | P90 | Max | +| --- | ---: | ---: | ---: | ---: | ---: | +| All 124 benchmarked shapes (diagnostic scope) | 1.0938x | 2.1825x | 2.1942x | 2.8423x | 3.8773x | + +#### Modeled after-init amortized synchronized E2E speedup + +| Public calls N | Min | Geomean | Median | P90 | Max | +| --- | ---: | ---: | ---: | ---: | ---: | +| 1 | 0.0865x | 1.2309x | 2.7840x | 3.3516x | 163.7986x | +| 10 | 0.1110x | 1.2589x | 2.6667x | 3.0908x | 137.3500x | +| 100 | 0.2850x | 1.5093x | 2.1682x | 2.9894x | 53.5582x | +| 1000 | 1.0120x | 1.9957x | 2.0271x | 2.8424x | 9.5268x | + +#### Modeled including-init amortized synchronized E2E speedup + +| Public calls N | Min | Geomean | Median | P90 | Max | +| --- | ---: | ---: | ---: | ---: | ---: | +| 1 | 0.0772x | 0.9450x | 2.2810x | 12.9363x | 24.8909x | +| 10 | 0.0814x | 0.9591x | 2.2786x | 12.7562x | 24.5273x | +| 100 | 0.1093x | 1.0709x | 2.2565x | 10.9253x | 21.4483x | +| 1000 | 0.3510x | 1.4680x | 1.5030x | 6.0454x | 10.3790x | + +- All three tables report synchronized host E2E speedups as `baseline/candidate`. `Hot steady-state` measures a repeated public call at each lane's declared hot cache state; its per-shape values supply the official metric used by the separate publication-floor section. +- `After-init amortized(N) = (first_compute + (N-1) * hot_median) / N`; it excludes init. +- `Including-init amortized(N) = (init + first_compute + (N-1) * hot_median) / N`; it includes init. Each latency formula is evaluated separately for baseline and candidate, then reported as `baseline/candidate`. Both amortized scenarios are composed from measured components, not a directly timed N-call loop. A lane without explicit init uses `I=0`. +- Init scope: `runtime_init_plus_standalone_shared_preprocess_support_once_per_validation_shard_process_device_operator`; composition: `runtime_init_plus_shared_preprocess_support_each_standalone_lane`; baseline has explicit init: `yes`. +- Cache policy: `synchronize_and_clear_after_each_completed_shape`; resident multi-shape cache benchmarked: `no`; cold order: `deterministic_balanced_per_publication_contract_portfolio`; init order: `alternate_candidate_baseline_first_by_validation_shard_parity_then_shared_support` +- Lifecycle timing convention: all three lifecycle tables are synchronized host E2E. Init/first-call brackets are CUPTI timestamp host diagnostics; separately, the hot GPU-span diagnostic remains strict correlated CUPTI activity timing. + +- Measured: `2026-07-09T04:30:10+00:00` +- Full summary: [`VALIDATION.json`](VALIDATION.json); per-shape results: [`BENCHMARK_RESULTS.json`](BENCHMARK_RESULTS.json) diff --git a/cake_exports/kmeans/benchmarks/all_expected_routes.json b/cake_exports/kmeans/benchmarks/all_expected_routes.json new file mode 100644 index 00000000..11ab68fa --- /dev/null +++ b/cake_exports/kmeans/benchmarks/all_expected_routes.json @@ -0,0 +1,1142 @@ +[ + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_lowdim_e50c_v1:launch_for_eval", + "selected_route": "lowdim_e50c_v1", + "shape": "d895_expanded_heldout_neighborhood_d80_b2_n2176_k1024_d80" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "d895_expanded_guard_miss_fallback_d128_b1_n1408_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "d895_expanded_forced_fallback_d128_b1_n1664_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1:launch_for_eval", + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "d895_expanded_guard_boundary_d144_b1_n128_k256_d144" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1:launch_for_eval", + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "d895_expanded_tail_divisibility_d176_b1_n1152_k512_d176" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_v1:launch_for_eval", + "selected_route": "d192_paired_repeated_mma_v1", + "shape": "d895_expanded_guard_overlap_d192_b1_n1024_k768_d192" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_highd_splitd_6fcf_v1:launch_for_eval", + "selected_route": "highd_splitd_single_tile_6fcf_v1", + "shape": "d895_expanded_random_legal_d320_b1_n1280_k512_d320" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_highd_splitd_6fcf_v1:launch_for_eval", + "selected_route": "highd_splitd_single_tile_6fcf_v1", + "shape": "d895_expanded_request_specific_d448_b1_n1280_k1024_d448" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1:launch_for_eval", + "selected_route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "shape": "d895_expanded_request_specific_d512_splitk_b1_n512_k8192_d512" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_microdim_hybrid_9c0d_v1:launch_for_eval", + "selected_route": "microdim_hybrid_9c0d_v1", + "shape": "d16_paired_b2_n4096_k1024_d16" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "microdim_pad64_d64_direct_v1", + "shape": "d16_fallback_b3_n2432_k512_d16" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "microdim_pad64_d64_direct_v1", + "shape": "d16_small_b4_n1024_k512_d16" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_microdim_hybrid_9c0d_v1:launch_for_eval", + "selected_route": "microdim_hybrid_9c0d_v1", + "shape": "d16_large_b8_n32768_k512_d16" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_microdim_hybrid_9c0d_v1:launch_for_eval", + "selected_route": "microdim_hybrid_9c0d_v1", + "shape": "d32_hugek_b1_n512_k8192_d32" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_microdim_hybrid_9c0d_v1:launch_for_eval", + "selected_route": "microdim_hybrid_9c0d_v1", + "shape": "d32_paired_b2_n4096_k1024_d32" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "microdim_pad64_d64_direct_v1", + "shape": "d32_fallback_b3_n2432_k512_d32" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "microdim_pad64_d64_direct_v1", + "shape": "d32_small_b4_n1024_k512_d32" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1:launch_for_eval", + "selected_route": "d64_direct_single64_1p2gap_9f2a_v1", + "shape": "d64_paired_b2_n4096_k1024_d64" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1:launch_for_eval", + "selected_route": "d64_direct_single64_1p2gap_9f2a_v1", + "shape": "d64_fallback_b3_n2432_k512_d64" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1:launch_for_eval", + "selected_route": "d64_direct_single64_1p2gap_9f2a_v1", + "shape": "d64_small_b4_n1024_k512_d64" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1:launch_for_eval", + "selected_route": "d64_direct_single64_1p2gap_9f2a_v1", + "shape": "d64_large_b8_n32768_k512_d64" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_lowdim_e50c_v1:launch_for_eval", + "selected_route": "lowdim_e50c_v1", + "shape": "d80_paired_b1_n2048_k1024_d80" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_lowdim_e50c_v1:launch_for_eval", + "selected_route": "lowdim_e50c_v1", + "shape": "d80_small_b2_n1024_k512_d80" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_lowdim_e50c_v1:launch_for_eval", + "selected_route": "lowdim_e50c_v1", + "shape": "d96_paired_b1_n1536_k1024_d96" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_lowdim_e50c_v1:launch_for_eval", + "selected_route": "lowdim_e50c_v1", + "shape": "d96_small_b2_n896_k512_d96" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_lowdim_e50c_v1:launch_for_eval", + "selected_route": "lowdim_e50c_v1", + "shape": "d96_random_b3_n6144_k1280_d96" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_lowdim_e50c_v1:launch_for_eval", + "selected_route": "lowdim_e50c_v1", + "shape": "d96_fallback_b5_n1664_k256_d96" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "small_grid_single_tile_v10", + "shape": "boundary_b1_n128_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "fallback_tiny_hugek_b1_n128_k4096_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "small_grid_single_tile_v10", + "shape": "small_b1_n256_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "paired_min_even_b1_n256_k768_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "paired_low_n_hugek_b1_n256_k4096_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "small_grid_single_tile_v10", + "shape": "small_b1_n896_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "fallback_odd_kover_b1_n896_k768_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "fallback_hugek_b1_n896_k16384_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "boundary_b1_n1024_k768_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "paired_kover_b1_n1024_k2048_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "forced_fallback_b1_n1152_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "fallback_first_after_small_b1_n1152_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "boundary_b1_n1280_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "paired_even_b1_n1536_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d128_splitk_priority_575c_v1:launch_for_eval", + "selected_route": "d128_splitk_priority_575c_v1", + "shape": "paired_largek_b1_n2048_k4096_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "fallback_odd_b1_n2176_k4096_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "tail_odd_b1_n4224_k1024_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "fallback_large_odd_b1_n8320_k1280_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "paired_large_b1_n65536_k4096_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "paired_large_n_b1_n131072_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "paired_kover_b2_n256_k768_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "fallback_low_n_hugek_b2_n384_k8192_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "paired_even_b2_n512_k8192_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "small_grid_single_tile_v10", + "shape": "small_random_b2_n768_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "random_b2_n768_k3072_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "small_grid_single_tile_v10", + "shape": "smoke_b2_n1024_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "fallback_first_odd_b2_n1408_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "paired_even_tail_b2_n1792_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "tail_even_b2_n2048_k2048_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "random_legal_b2_n4736_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "fallback_large_b2_n65664_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "small_grid_single_tile_v10", + "shape": "small_low_b3_n256_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "paired_first_even_b3_n1536_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "fallback_mid_odd_b3_n2432_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "random_legal_b3_n2560_k1536_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "random_b3_n6144_k1280_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "small_grid_single_tile_v10", + "shape": "small_b4_n128_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "random_legal_b4_n640_k1024_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "small_grid_single_tile_v10", + "shape": "boundary_b4_n1024_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "fallback_odd_b4_n1664_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "paired_widek_b4_n3072_k3072_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "small_grid_single_tile_v10", + "shape": "random_legal_b5_n512_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "random_b5_n3456_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "small_grid_single_tile_v10", + "shape": "small_cap_b6_n1024_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "paired_large_b6_n65536_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "small_grid_single_tile_v10", + "shape": "small_boundary_b7_n896_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "random_b7_n1024_k4096_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "fallback_b8_n3712_k2048_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "paired_b8_n4096_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "mid_b8_n8192_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "random_legal_b9_n3840_k1792_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "small_grid_single_tile_v10", + "shape": "small_b12_n640_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "small_grid_single_tile_v10", + "shape": "small_b16_n1024_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "large_b16_n32768_k1024_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "paper_b32_n75776_k1024_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_b23d_v2:launch_for_eval", + "selected_route": "d160_padded_single_repeated_mma_v2", + "shape": "d160_paired_b1_n2048_k2048_d160" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_b23d_v2:launch_for_eval", + "selected_route": "d160_padded_single_repeated_mma_v2", + "shape": "d160_fallback_b2_n2432_k1024_d160" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_v1:launch_for_eval", + "selected_route": "d192_single_repeated_mma_v1", + "shape": "d192_fallback_b1_n2176_k1024_d192" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_v1:launch_for_eval", + "selected_route": "d192_single_repeated_mma_v1", + "shape": "d192_small_b2_n1024_k512_d192" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_v1:launch_for_eval", + "selected_route": "d192_paired_repeated_mma_v1", + "shape": "d192_paired_b2_n2048_k2048_d192" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_v1:launch_for_eval", + "selected_route": "d192_paired_repeated_mma_v1", + "shape": "d192_large_b4_n32768_k1024_d192" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_v1:launch_for_eval", + "selected_route": "d256_single_repeated_mma_v1", + "shape": "d256_hugek_b1_n512_k8192_d256" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_v1:launch_for_eval", + "selected_route": "d256_single_repeated_mma_v1", + "shape": "d256_small_b1_n1024_k512_d256" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_v1:launch_for_eval", + "selected_route": "d256_single_repeated_mma_v1", + "shape": "d256_paired_b1_n4096_k4096_d256" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_v1:launch_for_eval", + "selected_route": "d256_single_repeated_mma_v1", + "shape": "d256_fallback_b2_n2432_k2048_d256" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_highd_splitd_6fcf_v1:launch_for_eval", + "selected_route": "highd_splitd_single_tile_6fcf_v1", + "shape": "d320_paired_b1_n2048_k4096_d320" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_highd_splitd_6fcf_v1:launch_for_eval", + "selected_route": "highd_splitd_single_tile_6fcf_v1", + "shape": "d320_large_b2_n16384_k1024_d320" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1:launch_for_eval", + "selected_route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "shape": "d384_hugek_b1_n768_k8192_d384" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_highd_splitd_6fcf_v1:launch_for_eval", + "selected_route": "highd_splitd_single_tile_6fcf_v1", + "shape": "d384_small_b1_n896_k512_d384" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_highd_splitd_6fcf_v1:launch_for_eval", + "selected_route": "highd_splitd_single_tile_6fcf_v1", + "shape": "d384_paired_b1_n2048_k4096_d384" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_highd_splitd_6fcf_v1:launch_for_eval", + "selected_route": "highd_splitd_single_tile_6fcf_v1", + "shape": "d384_fallback_b3_n3456_k1024_d384" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1:launch_for_eval", + "selected_route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "shape": "d448_hugek_b1_n512_k8192_d448" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1:launch_for_eval", + "selected_route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "shape": "d448_paired_b1_n2048_k4096_d448" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1:launch_for_eval", + "selected_route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "shape": "d512_hugek_b1_n512_k8192_d512" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_highd_splitd_6fcf_v1:launch_for_eval", + "selected_route": "highd_splitd_single_tile_6fcf_v1", + "shape": "d512_small_b1_n1024_k512_d512" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1:launch_for_eval", + "selected_route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "shape": "d512_paired_b1_n2048_k4096_d512" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_highd_splitd_6fcf_v1:launch_for_eval", + "selected_route": "highd_splitd_single_tile_6fcf_v1", + "shape": "d512_fallback_b2_n2432_k2048_d512" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_microdim_hybrid_9c0d_v1:launch_for_eval", + "selected_route": "microdim_hybrid_9c0d_v1", + "shape": "post_d895_d16_b4_n32768_k1024_d16" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_microdim_pipeline4_08f9_v4:launch_for_eval", + "selected_route": "microdim_pipeline4_08f9_v4", + "shape": "post_d895_d16_b8_n65536_k512_d16" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_microdim_hybrid_9c0d_v1:launch_for_eval", + "selected_route": "microdim_hybrid_9c0d_v1", + "shape": "post_d895_d32_b4_n32768_k1024_d32" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_microdim_pipeline4_08f9_v4:launch_for_eval", + "selected_route": "microdim_pipeline4_08f9_v4", + "shape": "post_d895_d32_b8_n65536_k512_d32" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_3328_d48_small_boundary_b1_n256_k256_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d48_b1_n512_k8192_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_d9d5_d48_highk_heldout_b2_n640_k4096_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_a2f8_d48_k1024_boundary_b2_n768_k1024_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_8f09_d48_medium_tail_b2_n1792_k512_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d48_b2_n2048_k512_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_5600_d48_low_n_boundary_b3_n1536_k256_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_c44f_d48_midk_boundary_b3_n2688_k768_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_ef00_d48_midn_boundary_b3_n3456_k512_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_68cf_d48_boundary_b4_n2304_k512_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_1d49_d48_lowk_boundary_b4_n3968_k256_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d48_b4_n8192_k1024_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_9ca0_d48_guard_boundary_b6_n12288_k512_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1:launch_for_eval", + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "shape": "post_d895_d112_b1_n256_k256_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1:launch_for_eval", + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "shape": "adjacent_8f09_d112_small_highk_b1_n384_k4096_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1:launch_for_eval", + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "shape": "post_d895_d112_b1_n512_k8192_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_1d49_d112_highk_tail_b1_n1408_k4096_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1:launch_for_eval", + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "shape": "adjacent_3328_d112_highk_low_n_b2_n512_k8192_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1:launch_for_eval", + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "shape": "post_d895_d112_b2_n2048_k512_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_c44f_d112_odd_tail_b2_n3200_k768_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1:launch_for_eval", + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "shape": "adjacent_d9d5_d112_highk_request_b3_n768_k8192_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_9ca0_d112_tail_div_b3_n3840_k768_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_a2f8_d112_lowk_tail_b4_n3456_k256_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_ef00_d112_odd_tail_b4_n3712_k1024_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d112_b4_n8192_k1024_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_5600_d112_odd_tile_tail_b5_n2176_k512_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_68cf_d112_tail_b5_n2944_k512_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "post_d895_d128_paired_b2_n262144_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "post_d895_d128_fallback_b3_n1920_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_8f09_d128_fallback_neighbor_b4_n4480_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_ef00_d128_forced_fallback_b4_n7552_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "post_d895_d128_fallback_b5_n2176_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_d9d5_d128_guard_miss_fallback_b5_n6016_k1024_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_1d49_d128_forced_fallback_b5_n7296_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_c44f_d128_forced_fallback_b6_n6272_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_a2f8_d128_odd_fallback_b6_n8576_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "post_d895_d128_fallback_b7_n2432_k1024_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_68cf_d128_forced_fallback_b7_n6016_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_9ca0_d128_exact_guard_neighbor_b8_n8064_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "d128_even_near_floor_v10_repair", + "shape": "post_d895_d128_paired_b8_n8192_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "adjacent_3328_d128_exact_guard_k_neighbor_b8_n8192_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_5600_d128_fallback_neighbor_b8_n8320_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1:launch_for_eval", + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d144_b1_n256_k256_d144" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1:launch_for_eval", + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d144_b1_n512_k8192_d144" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1:launch_for_eval", + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d144_b2_n2048_k1024_d144" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1:launch_for_eval", + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d144_b4_n8192_k1024_d144" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1:launch_for_eval", + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d176_b1_n256_k256_d176" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1:launch_for_eval", + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d176_b1_n512_k8192_d176" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1:launch_for_eval", + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d176_b2_n2048_k1024_d176" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1:launch_for_eval", + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d176_b4_n8192_k1024_d176" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d224_b1_n256_k256_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d224_b1_n512_k8192_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d224_b2_n2048_k1024_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_d9d5_d224_random_midk_b2_n2944_k1280_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_9ca0_d224_guard_overlap_b2_n4096_k256_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_ef00_d224_midn_overlap_b2_n5376_k512_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_a2f8_d224_highn_overlap_b2_n6144_k768_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_5600_d224_k512_overlap_b3_n3072_k512_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_3328_d224_tail_div_b3_n3840_k512_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_68cf_d224_overlap_b3_n5120_k1024_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4:launch_for_eval", + "selected_route": "d224_tmem_abi_repair_d17c_v4", + "shape": "adjacent_8f09_d224_low_n_boundary_b4_n1536_k256_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_c44f_d224_overlap_b4_n4480_k512_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d224_b4_n8192_k1024_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_1d49_d224_highn_overlap_b5_n5632_k768_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "post_d895_d288_b1_n256_k256_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_9ca0_d288_highk_low_n_b1_n384_k4096_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "post_d895_d288_b1_n512_k8192_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_a2f8_d288_splitk_probe_b1_n640_k4096_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_d9d5_d288_heldout_low_n_b1_n896_k4096_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_ef00_d288_highk_heldout_b1_n1664_k4096_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_8f09_d288_k768_overlap_b1_n2560_k768_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_c44f_d288_highk_heldout_b2_n768_k4096_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_5600_d288_mid_highk_b2_n1024_k2048_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_68cf_d288_boundary_b2_n1920_k512_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "post_d895_d288_b2_n2048_k1024_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_1d49_d288_low_n_heldout_b3_n1152_k2048_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_3328_d288_guard_overlap_b4_n8192_k256_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "post_d895_d288_b4_n8192_k1024_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d352_b1_n256_k256_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d352_exactd_splitk_c95c_v2:launch_for_eval", + "selected_route": "d352_exactd_splitk_c95c_v2", + "shape": "adjacent_9ca0_d352_splitk_boundary_b1_n512_k8192_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d352_exactd_splitk_c95c_v2:launch_for_eval", + "selected_route": "d352_exactd_splitk_c95c_v2", + "shape": "post_d895_d352_b1_n512_k8192_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_c44f_d352_random_b1_n3328_k768_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d352_exactd_splitk_c95c_v2:launch_for_eval", + "selected_route": "d352_exactd_splitk_c95c_v2", + "shape": "adjacent_5600_d352_splitk_edge_b2_n768_k4096_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d352_exactd_splitk_c95c_v2:launch_for_eval", + "selected_route": "d352_exactd_splitk_c95c_v2", + "shape": "adjacent_1d49_d352_request_highk_b2_n1024_k8192_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d352_exactd_splitk_c95c_v2:launch_for_eval", + "selected_route": "d352_exactd_splitk_c95c_v2", + "shape": "post_d895_d352_b2_n2048_k1024_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d352_exactd_splitk_c95c_v2:launch_for_eval", + "selected_route": "d352_exactd_splitk_c95c_v2", + "shape": "adjacent_ef00_d352_request_low_n_b3_n896_k8192_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_3328_d352_random_legal_b3_n2048_k768_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_68cf_d352_tail_b3_n2816_k768_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_a2f8_d352_smallk_boundary_b4_n1024_k256_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_8f09_d352_smallk_large_n_b4_n4096_k256_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d352_b4_n8192_k1024_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_d9d5_d352_random_b5_n2304_k768_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d416_b1_n256_k256_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_3328_d416_highk_low_n_b1_n384_k8192_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d416_b1_n512_k8192_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_9ca0_d416_splitk_min_b1_n1024_k4096_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_ef00_d416_random_midk_b1_n2176_k768_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_d9d5_d416_highk_request_b2_n640_k8192_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d416_b2_n2048_k1024_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_68cf_d416_overlap_b2_n2304_k1024_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_a2f8_d416_random_b2_n2560_k768_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_c44f_d416_highk_request_b3_n512_k8192_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_8f09_d416_midk_tail_b3_n3456_k768_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_5600_d416_smallk_boundary_b4_n2048_k256_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_1d49_d416_random_b4_n3840_k512_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d416_b4_n8192_k1024_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_3328_d480_min_boundary_b1_n128_k256_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d480_b1_n256_k256_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d480_splitk_k1024_eac2_v1:launch_for_eval", + "selected_route": "d480_splitk_k1024_eac2_v1", + "shape": "post_d895_d480_b1_n512_k8192_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d480_splitk_k1024_eac2_v1:launch_for_eval", + "selected_route": "d480_splitk_k1024_eac2_v1", + "shape": "adjacent_a2f8_d480_highk_request_b1_n896_k4096_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d480_splitk_k1024_eac2_v1:launch_for_eval", + "selected_route": "d480_splitk_k1024_eac2_v1", + "shape": "adjacent_5600_d480_random_b1_n1536_k1024_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d480_splitk_k1024_eac2_v1:launch_for_eval", + "selected_route": "d480_splitk_k1024_eac2_v1", + "shape": "adjacent_8f09_d480_highk_low_n_b2_n640_k4096_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d480_splitk_k1024_eac2_v1:launch_for_eval", + "selected_route": "d480_splitk_k1024_eac2_v1", + "shape": "post_d895_d480_b2_n2048_k1024_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_1d49_d480_smallk_boundary_b2_n2816_k256_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_9ca0_d480_splitd_large_n_b2_n4096_k512_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d480_splitk_k1024_eac2_v1:launch_for_eval", + "selected_route": "d480_splitk_k1024_eac2_v1", + "shape": "adjacent_d9d5_d480_heldout_highk_b3_n1024_k4096_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_ef00_d480_lowk_boundary_b3_n3200_k256_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_68cf_d480_boundary_b4_n1664_k512_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d480_b4_n8192_k1024_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_c44f_d480_boundary_b5_n2048_k512_d480" + } +] diff --git a/cake_exports/kmeans/benchmarks/benchmark.py b/cake_exports/kmeans/benchmarks/benchmark.py new file mode 100644 index 00000000..df81bdcd --- /dev/null +++ b/cake_exports/kmeans/benchmarks/benchmark.py @@ -0,0 +1,1656 @@ +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +import sys +import uuid +from pathlib import Path +from typing import Any + +ROOT = Path(__file__).resolve().parents[1] +SRC = ROOT / "src" +if str(SRC) not in sys.path: + sys.path.insert(0, str(SRC)) + +from flash_kmeans_shapes import ( # noqa: E402 + FLASH_KMEANS_EVOLUTION_ARTIFACT, + FLASH_KMEANS_EVOLUTION_SUMMARY, +) +from flash_kmeans_triton_h200 import euclid_assign_triton_h200 # noqa: E402 +from flash_kmeans_triton_h200_raw_adapter import ( # noqa: E402 + BASELINE_COMMIT, + PREPROCESS_IMPL, + TritonH20007cfRawAdapter, +) +from flash_kmeans_triton_h200_raw_adapter import ( # noqa: E402 + BASELINE_NAME as PUBLIC_RAW_BASELINE_NAME, +) +from flashlib_cake_kmeans._benchmark import ( # noqa: E402 + BenchResult, + HostCallTiming, + bench_gpu_time, + compare_runtime_lifecycles, + measure_host_call, + require_cupti, + runtime_lifecycle_metrics, +) + + +def _sum_host_timings(*timings: HostCallTiming) -> HostCallTiming: + """Compose sequential one-time setup costs for standalone-lane modeling.""" + + if not timings: + raise ValueError("at least one host timing is required") + return HostCallTiming( + host_enqueue_ms=sum(float(timing.host_enqueue_ms) for timing in timings), + synchronized_e2e_ms=sum(float(timing.synchronized_e2e_ms) for timing in timings), + ) + + +SHAPE_RECORDS = json.loads((Path(__file__).with_name("shape_records.json")).read_text(encoding="utf-8")) +ROUTE_MANIFEST = json.loads((Path(__file__).with_name("expected_routes.json")).read_text(encoding="utf-8")) +ALL_ROUTE_MANIFEST = json.loads((Path(__file__).with_name("all_expected_routes.json")).read_text(encoding="utf-8")) +EXPECTED_ROUTES = {row["shape"]: row["selected_route"] for row in ALL_ROUTE_MANIFEST} + + +def _shape_from_record(record: dict[str, Any]) -> dict[str, Any]: + params = dict(record["params"]) + recorded = dict(record.get("recorded", {})) + row = {"label": record["label"], **params} + row["runtime_coverage"] = bool(record.get("runtime_coverage", False)) + row["route"] = recorded.get("evolution_route", params.get("route")) + for key in ( + "evolution_kernel_ms", + "evolution_flashlib_ms", + "evolution_tflops", + "evolution_flashlib_equiv_tflops", + "evolution_speedup", + ): + row[key] = recorded.get(key, params.get(key)) + return row + + +FLASH_KMEANS_SHAPES = [_shape_from_record(record) for record in SHAPE_RECORDS] +_SHAPES_BY_LABEL = {row["label"]: row for row in FLASH_KMEANS_SHAPES} +FLASH_KMEANS_REGISTRY_SHAPES = [_SHAPES_BY_LABEL[route["shape"]] for route in ROUTE_MANIFEST] +SEMANTIC_ENTRYPOINT = "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval" +PRECOMPUTED_BASELINE_NAME = "triton_h200_07cf_precomputed" +BASELINE_REGISTRY_KEY = "triton_h200_07cf_dual_lane_v1" +# Updated together with benchmark_data.json by the registry sync gate. +BASELINE_REGISTRY_SHA256 = "bdfd30338aa614f09817af1498c1115a0b1732a3340a1acbc172d0e4dd4674c4" +REGISTRY_CANDIDATE_ENTRYPOINT = SEMANTIC_ENTRYPOINT +MEASURED_CANDIDATE_ENTRYPOINT = "flashlib_cake_kmeans.interface:FlashKMeansAssignRuntime.compute" +BASELINE_ENTRYPOINT = "benchmarks.flash_kmeans_triton_h200_raw_adapter:TritonH20007cfRawAdapter.compute" +PARITY_CANDIDATE_ENTRYPOINT = "flashlib_cake_kmeans.interface:flash_kmeans_assign_prepared" +PARITY_BASELINE_ENTRYPOINT = "benchmarks.flash_kmeans_triton_h200:euclid_assign_triton_h200" +REGISTRY_CANDIDATE_TIMING_BOUNDARY = ( + "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" +) +CANDIDATE_TIMING_BOUNDARY = ( + "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion" +) +BASELINE_TIMING_BOUNDARY = "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e" +PARITY_CANDIDATE_TIMING_BOUNDARY = "precomputed_norms_preallocated_output_prepared_assignment_gpu_span" +PARITY_BASELINE_TIMING_BOUNDARY = "precomputed_norms_preallocated_output_pinned_07cf_assignment_gpu_span" +BASELINE_REGISTRY_PROFILE = { + "registry_candidate_entrypoint": REGISTRY_CANDIDATE_ENTRYPOINT, + "shared_preprocess": { + "implementation_entrypoint": PREPROCESS_IMPL, + "source_sha256": "aa67813cf1cc39b8ae96970a737e926f7b3a65dac63dbeb6362f2dacf066e26e", + "result_source_sha256_field": "preprocess_source_sha256", + }, + "official": { + "role": "publication", + "candidate_entrypoint": MEASURED_CANDIDATE_ENTRYPOINT, + "baseline_name": PUBLIC_RAW_BASELINE_NAME, + "baseline_commit": BASELINE_COMMIT, + "baseline_entrypoint": BASELINE_ENTRYPOINT, + # This digest-bound identity belongs to the frozen benchmark registry. + # The live raw payload carries the request-owned official boundary. + "candidate_timing_boundary": REGISTRY_CANDIDATE_TIMING_BOUNDARY, + "baseline_timing_boundary": BASELINE_TIMING_BOUNDARY, + "candidate_timing_backend_field": "candidate_public_raw_timing_backend", + "baseline_timing_backend_field": "baseline_07cf_adapter_timing_backend", + "speedup_metric": "public_raw_e2e_speedup_vs_07cf_adapter", + "speedup_numerator_metric": "baseline_07cf_adapter_synchronized_e2e_ms", + "speedup_denominator_metric": "candidate_public_raw_synchronized_e2e_ms", + "timing_backend": "cupti", + }, + "parity": { + "role": "diagnostic_only", + "candidate_entrypoint": PARITY_CANDIDATE_ENTRYPOINT, + "baseline_name": PRECOMPUTED_BASELINE_NAME, + "baseline_commit": BASELINE_COMMIT, + "baseline_entrypoint": PARITY_BASELINE_ENTRYPOINT, + "candidate_timing_boundary": PARITY_CANDIDATE_TIMING_BOUNDARY, + "baseline_timing_boundary": PARITY_BASELINE_TIMING_BOUNDARY, + "candidate_timing_backend_field": "candidate_precomputed_timing_backend", + "baseline_timing_backend_field": "baseline_07cf_precomputed_timing_backend", + "speedup_metric": "precomputed_gpu_speedup_vs_07cf", + "speedup_numerator_metric": "baseline_07cf_precomputed_gpu_span_ms", + "speedup_denominator_metric": "candidate_precomputed_gpu_span_ms", + "timing_backend": "cupti", + }, +} +WARMUP_MS = 20.0 +BENCH_MS = 100.0 +DEFAULT_MEASUREMENT_ORDER_SEED = "flashlib-kmeans-export-paired-replay-v1" +MEASUREMENT_ROLES = ( + "candidate_public_raw", + "baseline_public_raw", + "candidate_precomputed", + "baseline_precomputed", +) + + +def _preprocess_source_digest() -> str: + """Hash the exact Python/CUDA pair-row-norm implementation in this export.""" + + source_dir = Path(__file__).resolve().parent + candidates = ( + ( + ROOT / "src/flashlib_cake_kmeans/_row_norm.py", + ROOT / "src/flashlib_cake_kmeans/_row_norm.cu", + ), + (source_dir / "row_norm.py", source_dir / "row_norm.cu"), + ) + for paths in candidates: + if all(path.is_file() for path in paths): + digest = hashlib.sha256() + for canonical_name, path in zip(("_row_norm.py", "_row_norm.cu"), paths, strict=True): + digest.update(canonical_name.encode("utf-8")) + digest.update(b"\0") + digest.update(path.read_bytes()) + digest.update(b"\0") + return digest.hexdigest() + searched = ", ".join(str(path) for pair in candidates for path in pair) + raise RuntimeError(f"cannot locate exported pair-row-norm sources; searched {searched}") + + +def _timing_diagnostics(timing: Any, *, primary_metric: str) -> dict[str, Any]: + return { + "lane_primary_metric": primary_metric, + # KMeans uses an adaptive 100 ms window that can contain thousands of + # iterations. The generated BenchResult retains every exact sample; + # the publication JSON records medians and sample count to avoid + # multiplying artifact size across the 228-shape portfolio. + "sample_count": len(timing.times_ms), + "gpu_span_ms": {"median": timing.median_gpu_span_ms}, + "kernel_sum_ms": {"median": timing.median_kernel_sum_ms}, + "active_union_ms": {"median": timing.median_active_union_ms}, + "inter_kernel_gap_ms": {"median": timing.median_inter_kernel_gap_ms}, + "activity_count": {"median": timing.median_activity_count}, + "correlated_launch_activity_count": {"median": timing.median_launch_activity_count}, + "correlated_kernel_activity_count": {"median": timing.median_kernel_activity_count}, + "submission_ms": {"median": timing.median_submission_ms}, + "host_enqueue_ms": {"median": timing.median_host_enqueue_ms}, + "synchronized_e2e_ms": {"median": timing.median_synchronized_e2e_ms}, + "cold_first_call": { + "host_enqueue_ms": timing.cold_first_call_host_enqueue_ms, + "synchronized_e2e_ms": timing.cold_first_call_synchronized_e2e_ms, + }, + } + + +def _pair_order( + label: str, + pair_name: str, + pair_index: int, + order_seed: str, +) -> tuple[str, str]: + pairs = { + "public": ("candidate_public_raw", "baseline_public_raw"), + "precomputed": ("candidate_precomputed", "baseline_precomputed"), + } + pair = pairs[pair_name] + digest = hashlib.sha256( + f"{order_seed}:{label}:{pair_name}:{pair_index}".encode() + ).digest() + return pair if digest[0] & 1 else tuple(reversed(pair)) + + +def _interleaved_pair_schedule( + label: str, + *, + public_pair_count: int, + precomputed_pair_count: int, + order_seed: str, +) -> tuple[str, ...]: + """Build balanced adjacent A/B pairs for one shared CUPTI session.""" + + if not isinstance(order_seed, str) or not order_seed: + raise ValueError("measurement order seed must be a non-empty string") + if public_pair_count <= 0 or precomputed_pair_count <= 0: + raise ValueError("interleaved pair counts must be positive") + pair_slots = [ + ((index + 0.5) / public_pair_count, "public", index) + for index in range(public_pair_count) + ] + pair_slots.extend( + ((index + 0.5) / precomputed_pair_count, "precomputed", index) + for index in range(precomputed_pair_count) + ) + pair_slots.sort( + key=lambda item: ( + item[0], + hashlib.sha256( + f"{order_seed}:{label}:pair-slot:{item[1]}:{item[2]}".encode() + ).digest(), + ) + ) + schedule: list[str] = [] + for _position, pair_name, pair_index in pair_slots: + schedule.extend(_pair_order(label, pair_name, pair_index, order_seed)) + counts = {role: schedule.count(role) for role in MEASUREMENT_ROLES} + expected = { + "candidate_public_raw": public_pair_count, + "baseline_public_raw": public_pair_count, + "candidate_precomputed": precomputed_pair_count, + "baseline_precomputed": precomputed_pair_count, + } + if counts != expected: + raise RuntimeError(f"interleaved role schedule is incomplete: {counts!r} != {expected!r}") + return tuple(schedule) + + +def _slice_bench_result( + result: BenchResult, + indices: list[int], + *, + cold_first_call: HostCallTiming | None, +) -> BenchResult: + if not indices: + raise ValueError("each paired role requires at least one reportable sample") + + def select(values): + if values is None: + return None + if len(values) != len(result.times_ms): + raise RuntimeError("shared CUPTI result fields have inconsistent sample counts") + return [values[index] for index in indices] + + return BenchResult( + times_ms=select(result.times_ms), + backend=result.backend, + kernel_sum_times_ms=select(result.kernel_sum_times_ms), + inter_kernel_gap_times_ms=select(result.inter_kernel_gap_times_ms), + active_union_times_ms=select(result.active_union_times_ms), + activity_counts=select(result.activity_counts), + launch_activity_counts=select(result.launch_activity_counts), + kernel_activity_counts=select(result.kernel_activity_counts), + submission_times_ms=select(result.submission_times_ms), + synchronized_e2e_times_ms=select(result.synchronized_e2e_times_ms), + cold_first_call_host_enqueue_ms=( + cold_first_call.host_enqueue_ms if cold_first_call is not None else None + ), + cold_first_call_synchronized_e2e_ms=( + cold_first_call.synchronized_e2e_ms if cold_first_call is not None else None + ), + ) + + +def _bench_interleaved_roles( + *, + label: str, + order_seed: str, + role_factories: dict[str, Any], + cold_first_calls: dict[str, HostCallTiming | None], +) -> tuple[dict[str, BenchResult], dict[str, Any]]: + """Measure all reportable roles in one interleaved CUPTI activity session.""" + + if tuple(role_factories) != MEASUREMENT_ROLES: + raise ValueError("paired benchmark role factories are incomplete or out of order") + probes: dict[str, BenchResult] = {} + for role in MEASUREMENT_ROLES: + probes[role] = bench_gpu_time( + role_factories[role](), + warmup_iters=5, + bench_iters=20, + cold_l2=True, + ) + estimates = {role: float(probes[role].median_ms) for role in MEASUREMENT_ROLES} + if any(not math.isfinite(value) or value <= 0.0 for value in estimates.values()): + raise RuntimeError(f"paired benchmark probes produced invalid estimates: {estimates!r}") + + public_bench_pairs = max( + 8, + math.ceil(BENCH_MS / estimates["candidate_public_raw"]), + math.ceil(BENCH_MS / estimates["baseline_public_raw"]), + ) + precomputed_bench_pairs = max( + 8, + math.ceil(BENCH_MS / estimates["candidate_precomputed"]), + math.ceil(BENCH_MS / estimates["baseline_precomputed"]), + ) + public_warmup_pairs = max( + 1, + math.ceil(WARMUP_MS / estimates["candidate_public_raw"]), + math.ceil(WARMUP_MS / estimates["baseline_public_raw"]), + ) + precomputed_warmup_pairs = max( + 1, + math.ceil(WARMUP_MS / estimates["candidate_precomputed"]), + math.ceil(WARMUP_MS / estimates["baseline_precomputed"]), + ) + warmup_schedule = _interleaved_pair_schedule( + label, + public_pair_count=public_warmup_pairs, + precomputed_pair_count=precomputed_warmup_pairs, + order_seed=f"{order_seed}:warmup", + ) + measurement_schedule = _interleaved_pair_schedule( + label, + public_pair_count=public_bench_pairs, + precomputed_pair_count=precomputed_bench_pairs, + order_seed=order_seed, + ) + combined_schedule = (*warmup_schedule, *measurement_schedule) + role_callables = {role: role_factories[role]() for role in MEASUREMENT_ROLES} + next_index = 0 + + def invoke_scheduled_role(): + nonlocal next_index + role = combined_schedule[next_index] + next_index += 1 + return role_callables[role]() + + shared = bench_gpu_time( + invoke_scheduled_role, + warmup_iters=len(warmup_schedule), + bench_iters=len(measurement_schedule), + cold_l2=True, + ) + if next_index != len(combined_schedule): + raise RuntimeError("paired benchmark did not consume its complete role schedule") + timings: dict[str, BenchResult] = {} + sample_counts: dict[str, int] = {} + for role in MEASUREMENT_ROLES: + indices = [index for index, scheduled_role in enumerate(measurement_schedule) if scheduled_role == role] + timings[role] = _slice_bench_result( + shared, + indices, + cold_first_call=cold_first_calls.get(role), + ) + sample_counts[role] = len(indices) + + schedule_text = "\n".join(measurement_schedule).encode() + provenance = { + "schema": "flash-kmeans-interleaved-cupti-session-v1", + "same_cupti_session": True, + "interleaved": True, + "adjacent_candidate_baseline_pairs": True, + "reportable_cupti_session_count": 1, + "adaptive_probe_cupti_session_count": len(MEASUREMENT_ROLES), + "adaptive_probes_reportable": False, + "order_seed": order_seed, + "measurement_schedule_sha256": hashlib.sha256(schedule_text).hexdigest(), + "measurement_schedule_iteration_count": len(measurement_schedule), + "warmup_schedule_iteration_count": len(warmup_schedule), + "role_sample_counts": sample_counts, + "public_pair_count": public_bench_pairs, + "precomputed_pair_count": precomputed_bench_pairs, + "first_reportable_roles": list(measurement_schedule[:16]), + } + return timings, provenance + + +def _label_seed(label: str) -> int: + digest = hashlib.sha256(label.encode("utf-8")).digest() + return int.from_bytes(digest[:8], "little") % (2**31) + + +_REGISTRY_COLD_LABELS = frozenset(row["label"] for row in FLASH_KMEANS_REGISTRY_SHAPES) +_CANDIDATE_FIRST_COLD_LABELS_BY_SEED: dict[str, frozenset[str]] = {} + + +def _candidate_first_cold_labels(order_seed: str) -> frozenset[str]: + cached = _CANDIDATE_FIRST_COLD_LABELS_BY_SEED.get(order_seed) + if cached is None: + ordered = sorted( + (row["label"] for row in FLASH_KMEANS_REGISTRY_SHAPES), + key=lambda label: hashlib.sha256(f"{order_seed}:cold:{label}".encode()).digest(), + ) + cached = frozenset(ordered[::2]) + _CANDIDATE_FIRST_COLD_LABELS_BY_SEED[order_seed] = cached + return cached + + +def _cold_measurement_order(label: str, order_seed: str) -> tuple[str, str]: + """Counterbalance process-shared cold effects across contract shapes. + + Rank-alternate over the seeded hash order of the full contract portfolio + so the split is exactly balanced (|candidate_first - baseline_first| <= 1) + for every shard decomposition, matching the declared + ``deterministic_balanced_per_publication_contract_portfolio`` policy. + """ + + if label in _REGISTRY_COLD_LABELS: + first_candidate = label in _candidate_first_cold_labels(order_seed) + else: + # Runtime-coverage diagnostic rows sit outside the publication + # contract balance gate; keep the per-label seeded coin flip. + digest = hashlib.sha256(f"{order_seed}:cold:{label}".encode()).digest() + first_candidate = bool(digest[0] & 1) + return ("candidate", "baseline") if first_candidate else ("baseline", "candidate") + + +def _alternating_call(first, second): + """Return a callable that alternates two equivalent fresh-pointer cases.""" + + calls = (first, second) + next_index = 0 + + def invoke(): + nonlocal next_index + current = next_index + next_index ^= 1 + return calls[current]() + + return invoke + + +def _required_median(value: float | None, *, metric: str) -> float: + if value is None: + raise RuntimeError(f"CUPTI benchmark did not report required {metric}") + return float(value) + + +def _dual_lane_speedups( + *, + candidate_public_raw_timing: Any, + baseline_public_raw_timing: Any, + candidate_precomputed_timing: Any, + baseline_precomputed_timing: Any, +) -> dict[str, float]: + """Return the one official E2E ratio and the precomputed GPU parity ratio.""" + + candidate_public_raw_e2e = _required_median( + candidate_public_raw_timing.median_synchronized_e2e_ms, + metric="candidate public-raw synchronized_e2e_ms", + ) + baseline_public_raw_e2e = _required_median( + baseline_public_raw_timing.median_synchronized_e2e_ms, + metric="07cf raw-adapter synchronized_e2e_ms", + ) + return { + "public_raw_e2e_speedup_vs_07cf_adapter": baseline_public_raw_e2e / candidate_public_raw_e2e, + "precomputed_gpu_speedup_vs_07cf": ( + float(baseline_precomputed_timing.median_gpu_span_ms) + / float(candidate_precomputed_timing.median_gpu_span_ms) + ), + } + + +def _public_prepared_ratios(public_timing: Any, prepared_timing: Any) -> dict[str, float]: + """Return public/prepared ratios from one per-shape measurement session.""" + + public_submission = _required_median( + public_timing.median_submission_ms, + metric="public submission_ms", + ) + prepared_submission = _required_median( + prepared_timing.median_submission_ms, + metric="prepared submission_ms", + ) + public_e2e = _required_median( + public_timing.median_synchronized_e2e_ms, + metric="public synchronized_e2e_ms", + ) + prepared_e2e = _required_median( + prepared_timing.median_synchronized_e2e_ms, + metric="prepared synchronized_e2e_ms", + ) + denominators = { + "gpu_span_ms": float(prepared_timing.median_gpu_span_ms), + "submission_ms": prepared_submission, + "synchronized_e2e_ms": prepared_e2e, + } + if any(value <= 0.0 for value in denominators.values()): + raise RuntimeError(f"prepared timing ratios require positive denominators, got {denominators!r}") + return { + "gpu_span": float(public_timing.median_gpu_span_ms) / denominators["gpu_span_ms"], + "submission": public_submission / denominators["submission_ms"], + "synchronized_e2e": public_e2e / denominators["synchronized_e2e_ms"], + } + + +def _cold_call_fields(timing: Any | None) -> dict[str, float | None]: + return { + "submission_ms": timing.host_enqueue_ms if timing is not None else None, + "host_enqueue_ms": timing.host_enqueue_ms if timing is not None else None, + "synchronized_e2e_ms": timing.synchronized_e2e_ms if timing is not None else None, + } + + +def _shape_key(row: dict[str, Any]) -> tuple[int, int, int, int, str]: + return ( + int(row["B"]), + int(row["N"]), + int(row["D"]), + int(row["K"]), + str(row.get("dtype", "bfloat16")), + ) + + +def _selected_rows(args: argparse.Namespace) -> list[dict[str, Any]]: + rows = list(FLASH_KMEANS_SHAPES if args.include_runtime_coverage else FLASH_KMEANS_REGISTRY_SHAPES) + if args.source: + wanted_sources = set(args.source) + rows = [row for row in rows if row["source"] in wanted_sources] + if args.shape: + by_label = {row["label"]: row for row in rows} + missing = sorted(set(args.shape) - set(by_label)) + if missing: + available = ", ".join(sorted(by_label)) + raise SystemExit(f"unknown shape label(s) {missing}. Available: {available}") + rows = [by_label[label] for label in args.shape] + if args.unique: + seen: set[tuple[int, int, int, int, str]] = set() + unique_rows: list[dict[str, Any]] = [] + for row in rows: + key = _shape_key(row) + if key in seen: + continue + seen.add(key) + unique_rows.append(row) + rows = unique_rows + if args.limit is not None: + rows = rows[: args.limit] + rows = rows[args.shard_index :: args.shard_count] + return rows + + +def _write_json_atomic(path: Path, text: str) -> None: + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") + temporary.write_text(text + "\n", encoding="utf-8") + temporary.replace(path) + + +def _make_inputs(row: dict[str, Any], *, variant: int = 0): + import torch + + generator = torch.Generator(device="cuda") + seed = int(row.get("seed", _label_seed(str(row["label"])))) + if variant != 0: + seed = _label_seed(f"{row['label']}:fresh-pointer:{variant}:{seed}") + generator.manual_seed(seed) + x = torch.randn( + (int(row["B"]), int(row["N"]), int(row["D"])), + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ).contiguous() + centroids = torch.randn( + (int(row["B"]), int(row["K"]), int(row["D"])), + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ).contiguous() + return x, centroids + + +def _measure_shared_preprocess_cold_compile(row: dict[str, Any], *, arch: str): + """Compile the shared row-norm support once, outside either comparison lane.""" + + import torch + from flashlib_cake_kmeans._row_norm import prepare_bf16_pair_row_norm + + x, centroids = _make_inputs(row, variant=0xC01D) + x_sq = torch.empty((int(row["B"]), int(row["N"])), dtype=torch.float32, device=x.device) + c_sq = torch.empty((int(row["B"]), int(row["K"])), dtype=torch.float32, device=x.device) + # Fixture generation is asynchronous. Complete it before the cold bracket + # so neither lane inherits unrelated random-generation/allocation work. + torch.cuda.synchronize() + + def prepare_shared_support(): + prepared = prepare_bf16_pair_row_norm( + x, + centroids, + x_sq, + c_sq, + compute_x=True, + compute_c=True, + arch=arch, + ) + prepared.release_bound_callers(x_sq) + + _, cold_compile = measure_host_call(prepare_shared_support) + return cold_compile + + +def _reference_assign(x, centroids, *, chunk_rows: int): + import torch + + bsz = int(x.shape[0]) + n_points = int(x.shape[1]) + x_f32 = x.float() + c_f32 = centroids.float() + c_sq = (c_f32 * c_f32).sum(-1) + ref = torch.empty((bsz, n_points), dtype=torch.int32, device=x.device) + with torch.no_grad(): + for b in range(bsz): + c_t = c_f32[b].transpose(0, 1).contiguous() + c_bias = 0.5 * c_sq[b] + for start in range(0, n_points, chunk_rows): + q = x_f32[b, start : start + chunk_rows] + scores = torch.matmul(q, c_t) - c_bias.unsqueeze(0) + ref[b, start : start + q.shape[0]] = scores.argmax(dim=-1).to(torch.int32) + return ref + + +def _assignment_correctness(cluster_ids, ref, x, centroids) -> dict[str, Any]: + """Validate assignment indices while accepting exact-distance ties.""" + + import torch + + matches = cluster_ids == ref + mismatch_count = int((~matches).sum().item()) + diagnostics: dict[str, Any] = { + "match_rate": float(matches.float().mean().item()), + "mismatch_count": mismatch_count, + } + if mismatch_count == 0: + diagnostics.update({"correct": True, "tie_inclusive": False}) + return diagnostics + + dim = int(x.shape[-1]) + pred_idx = cluster_ids.to(torch.int64).unsqueeze(-1).expand(-1, -1, dim) + ref_idx = ref.to(torch.int64).unsqueeze(-1).expand(-1, -1, dim) + pred_centroids = torch.gather(centroids, 1, pred_idx).float() + ref_centroids = torch.gather(centroids, 1, ref_idx).float() + points = x.float() + pred_dist = ((points - pred_centroids) ** 2).sum(-1) + ref_dist = ((points - ref_centroids) ** 2).sum(-1) + distance_delta = (pred_dist - ref_dist).abs() + tie_ok = matches | (distance_delta <= 1.0e-3) + diagnostics.update( + { + "correct": bool(tie_ok.all().item()), + "tie_inclusive": True, + "tie_inclusive_match_rate": float(tie_ok.float().mean().item()), + "max_selected_distance_delta": float(distance_delta.max().item()), + } + ) + return diagnostics + + +def _run_shape( + row: dict[str, Any], + *, + runtime: Any, + baseline_adapter: Any, + candidate_init_timing: Any = None, + baseline_init_timing: Any = None, + arch: str | None, + correctness: bool, + benchmark: bool, + reference_chunk_rows: int, + measurement_session_id: str, + measurement_order_seed: str, +) -> dict[str, Any]: + import torch + from flashlib_cake_kmeans import ( + flash_kmeans_assign_prepared, + prepare_flash_kmeans_assign, + ) + + x, centroids = _make_inputs(row, variant=0) + fresh_x, fresh_centroids = _make_inputs(row, variant=1) + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + fresh_x_sq = (fresh_x.float() ** 2).sum(-1).contiguous() + fresh_c_sq = (fresh_centroids.float() ** 2).sum(-1).contiguous() + out_shape = (int(row["B"]), int(row["N"])) + candidate_precomputed_out = torch.empty(out_shape, dtype=torch.int32, device=x.device) + fresh_candidate_precomputed_out = torch.empty_like(candidate_precomputed_out) + baseline_precomputed_out = torch.empty_like(candidate_precomputed_out) + fresh_baseline_precomputed_out = torch.empty_like(candidate_precomputed_out) + + pointer_pairs = ( + (x, fresh_x), + (centroids, fresh_centroids), + (x_sq, fresh_x_sq), + (c_sq, fresh_c_sq), + ) + if any(first.data_ptr() == second.data_ptr() for first, second in pointer_pairs): + raise RuntimeError("fresh-pointer benchmark inputs unexpectedly alias the first input set") + # Keep cold diagnostics honest: random input generation, explicit parity + # norms, and preallocated parity outputs all precede every measured call. + torch.cuda.synchronize() + + def baseline_public_raw_first_call(): + return baseline_adapter.compute(x, centroids, return_info=True) + + def baseline_public_raw_fresh_pointer_call(): + return baseline_adapter.compute(fresh_x, fresh_centroids, return_info=True) + + def baseline_precomputed_first_call(): + return euclid_assign_triton_h200( + x, + centroids, + x_sq, + c_sq, + out=baseline_precomputed_out, + ) + + def baseline_precomputed_fresh_pointer_call(): + return euclid_assign_triton_h200( + fresh_x, + fresh_centroids, + fresh_x_sq, + fresh_c_sq, + out=fresh_baseline_precomputed_out, + ) + + def candidate_public_raw_shape_first_call(): + return runtime.compute( + x, + centroids, + return_info=True, + ) + + def candidate_public_raw_fresh_pointer_call(): + return runtime.compute( + fresh_x, + fresh_centroids, + return_info=True, + ) + + baseline_public_raw_shape_cold_call = None + baseline_public_raw_fresh_cold_call = None + baseline_precomputed_first_cold_call = None + baseline_precomputed_fresh_cold_call = None + candidate_public_raw_shape_cold_call = None + candidate_public_raw_fresh_cold_call = None + candidate_precomputed_prepare_cold_call = None + fresh_candidate_precomputed_prepare_cold_call = None + candidate_precomputed_first_cold_call = None + fresh_candidate_precomputed_first_cold_call = None + cold_measurement_order = _cold_measurement_order( + str(row["label"]), + measurement_order_seed, + ) + if benchmark: + for lane in cold_measurement_order: + if lane == "baseline": + ( + (baseline_public_raw_out, baseline_public_raw_shape_info), + baseline_public_raw_shape_cold_call, + ) = measure_host_call(baseline_public_raw_first_call) + else: + ( + (candidate_public_raw_out, candidate_public_raw_shape_info), + candidate_public_raw_shape_cold_call, + ) = measure_host_call(candidate_public_raw_shape_first_call) + for lane in cold_measurement_order: + if lane == "baseline": + ( + (fresh_baseline_public_raw_out, baseline_public_raw_fresh_info), + baseline_public_raw_fresh_cold_call, + ) = measure_host_call(baseline_public_raw_fresh_pointer_call) + else: + ( + (fresh_candidate_public_raw_out, candidate_public_raw_info), + candidate_public_raw_fresh_cold_call, + ) = measure_host_call(candidate_public_raw_fresh_pointer_call) + candidate_prepared, candidate_precomputed_prepare_cold_call = measure_host_call( + lambda: prepare_flash_kmeans_assign( + x, + centroids, + out=candidate_precomputed_out, + x_sq=x_sq, + c_sq=c_sq, + arch=arch, + ) + ) + fresh_candidate_prepared, fresh_candidate_precomputed_prepare_cold_call = measure_host_call( + lambda: prepare_flash_kmeans_assign( + fresh_x, + fresh_centroids, + out=fresh_candidate_precomputed_out, + x_sq=fresh_x_sq, + c_sq=fresh_c_sq, + arch=arch, + ) + ) + ( + ( + candidate_precomputed_result, + candidate_precomputed_info, + ), + candidate_precomputed_first_cold_call, + ) = measure_host_call(lambda: flash_kmeans_assign_prepared(candidate_prepared, return_info=True)) + ( + ( + fresh_candidate_precomputed_result, + fresh_candidate_precomputed_info, + ), + fresh_candidate_precomputed_first_cold_call, + ) = measure_host_call(lambda: flash_kmeans_assign_prepared(fresh_candidate_prepared, return_info=True)) + ( + ( + baseline_precomputed_result, + baseline_precomputed_config, + ), + baseline_precomputed_first_cold_call, + ) = measure_host_call(baseline_precomputed_first_call) + ( + ( + fresh_baseline_precomputed_result, + fresh_baseline_precomputed_config, + ), + baseline_precomputed_fresh_cold_call, + ) = measure_host_call(baseline_precomputed_fresh_pointer_call) + else: + ( + baseline_public_raw_out, + baseline_public_raw_shape_info, + ) = baseline_public_raw_first_call() + ( + fresh_baseline_public_raw_out, + baseline_public_raw_fresh_info, + ) = baseline_public_raw_fresh_pointer_call() + candidate_public_raw_out, candidate_public_raw_shape_info = candidate_public_raw_shape_first_call() + fresh_candidate_public_raw_out, candidate_public_raw_info = candidate_public_raw_fresh_pointer_call() + candidate_prepared = prepare_flash_kmeans_assign( + x, + centroids, + out=candidate_precomputed_out, + x_sq=x_sq, + c_sq=c_sq, + arch=arch, + ) + fresh_candidate_prepared = prepare_flash_kmeans_assign( + fresh_x, + fresh_centroids, + out=fresh_candidate_precomputed_out, + x_sq=fresh_x_sq, + c_sq=fresh_c_sq, + arch=arch, + ) + candidate_precomputed_result, candidate_precomputed_info = flash_kmeans_assign_prepared( + candidate_prepared, + return_info=True, + ) + fresh_candidate_precomputed_result, fresh_candidate_precomputed_info = flash_kmeans_assign_prepared( + fresh_candidate_prepared, + return_info=True, + ) + baseline_precomputed_result, baseline_precomputed_config = baseline_precomputed_first_call() + fresh_baseline_precomputed_result, fresh_baseline_precomputed_config = baseline_precomputed_fresh_pointer_call() + torch.cuda.synchronize() + + baseline_configs = ( + baseline_public_raw_shape_info["triton_h200_07cf_config"], + baseline_public_raw_fresh_info["triton_h200_07cf_config"], + baseline_precomputed_config, + fresh_baseline_precomputed_config, + ) + if any(config != baseline_configs[0] for config in baseline_configs[1:]): + raise RuntimeError("07cf Triton selected different configs across equivalent dual-lane calls") + if candidate_public_raw_info["selected_route"] != candidate_public_raw_shape_info["selected_route"]: + raise RuntimeError( + "runtime shape-first and fresh-pointer calls selected different routes: " + f"{candidate_public_raw_shape_info['selected_route']!r} != " + f"{candidate_public_raw_info['selected_route']!r}" + ) + if not candidate_public_raw_info["runtime_cache_hit"]: + raise RuntimeError("fresh-pointer runtime.compute call did not reuse the shape/stream launch plan") + if not baseline_public_raw_fresh_info["runtime_cache_hit"]: + raise RuntimeError("fresh-pointer 07cf raw-adapter call did not reuse the shape/stream norm plan") + if candidate_public_raw_info.get("hot_launch_path") != "cuda_graph": + raise RuntimeError( + "Flash-KMeans signature did not capture into a CUDA graph: " + f"hot_launch_path={candidate_public_raw_info.get('hot_launch_path')!r}, " + f"graph_capture_error={candidate_public_raw_info.get('graph_capture_error')!r}; " + "every contract signature freezes to a LaunchPlan today, so a " + "prepared-path fallback means the graph runtime silently degraded" + ) + candidate_shape_norm_fields = tuple(candidate_public_raw_shape_info.get("norm_compute_fields", ())) + candidate_fresh_norm_fields = tuple(candidate_public_raw_info.get("norm_compute_fields", ())) + if candidate_shape_norm_fields != candidate_fresh_norm_fields: + raise RuntimeError( + "runtime shape-first and fresh-pointer calls computed different norm fields: " + f"{candidate_shape_norm_fields!r} != {candidate_fresh_norm_fields!r}" + ) + for prepared_info in (candidate_precomputed_info, fresh_candidate_precomputed_info): + if prepared_info["selected_route"] != candidate_public_raw_info["selected_route"]: + raise RuntimeError( + "public-raw and precomputed Flash-KMeans calls selected different routes: " + f"{candidate_public_raw_info['selected_route']!r} != {prepared_info['selected_route']!r}" + ) + if candidate_precomputed_result is not candidate_precomputed_out: + raise RuntimeError("prepared candidate did not preserve its first caller-owned output") + if fresh_candidate_precomputed_result is not fresh_candidate_precomputed_out: + raise RuntimeError("prepared candidate did not preserve its fresh-pointer caller-owned output") + + prepared_direct_launcher = candidate_prepared.launch_plan.direct_launcher + fresh_prepared_direct_launcher = fresh_candidate_prepared.launch_plan.direct_launcher + prepared_path_proof = { + "resolved_direct_launcher": bool(callable(prepared_direct_launcher)), + "stable_direct_launcher_tokens": [ + f"{type(prepared_direct_launcher).__module__}.{type(prepared_direct_launcher).__qualname__}" + f"@{id(prepared_direct_launcher):x}", + f"{type(fresh_prepared_direct_launcher).__module__}." + f"{type(fresh_prepared_direct_launcher).__qualname__}@{id(fresh_prepared_direct_launcher):x}", + ], + "caller_owned_outputs": True, + "scratch_reused": True, + "stream_bound": bool( + candidate_precomputed_info["stream_handle"] == candidate_prepared.launch_plan.stream_handle + and fresh_candidate_precomputed_info["stream_handle"] + == fresh_candidate_prepared.launch_plan.stream_handle + ), + "root_dispatch_traversal_count": 0, + "parent_dispatch_traversal_count": 0, + "internal_device_wide_synchronization_count": 0, + "proof_scope": "prepared_direct_launcher_identity_and_source_invariants", + } + if not all( + prepared_path_proof[key] + for key in ("resolved_direct_launcher", "caller_owned_outputs", "scratch_reused", "stream_bound") + ): + raise RuntimeError(f"prepared candidate hot-path proof failed: {prepared_path_proof!r}") + + candidate_shape_first_was_cache_hit = bool(candidate_public_raw_shape_info["runtime_cache_hit"]) + baseline_shape_first_was_cache_hit = bool(baseline_public_raw_shape_info["runtime_cache_hit"]) + + result: dict[str, Any] = { + "shape": row["label"], + "label": row["label"], + "source": row["source"], + "seed": row.get("seed"), + "runtime_coverage": bool(row.get("runtime_coverage", False)), + "B": int(row["B"]), + "N": int(row["N"]), + "D": int(row["D"]), + "K": int(row["K"]), + "dtype": row.get("dtype", "bfloat16"), + "semantic_entrypoint": candidate_public_raw_info["semantic_entrypoint"], + "expected_route": EXPECTED_ROUTES[row["label"]], + "selected_route": candidate_public_raw_info["selected_route"], + "child_route": candidate_public_raw_info.get("child_route"), + "route_matches_expected": bool(candidate_public_raw_info["selected_route"] == EXPECTED_ROUTES[row["label"]]), + "evolution_kernel_ms": row["evolution_kernel_ms"], + "evolution_flashlib_ms": row["evolution_flashlib_ms"], + "evolution_tflops": row["evolution_tflops"], + "evolution_speedup": row["evolution_speedup"], + "public_raw_baseline_name": PUBLIC_RAW_BASELINE_NAME, + "precomputed_baseline_name": PRECOMPUTED_BASELINE_NAME, + "baseline_commit": BASELINE_COMMIT, + "triton_h200_07cf_config": baseline_configs[0], + "measurement_session_id": measurement_session_id, + "candidate_public_raw_assignment_launch_count": candidate_public_raw_info["assignment_launch_count"], + "candidate_public_raw_norm_launch_count": candidate_public_raw_info["norm_launch_count"], + "candidate_public_raw_norm_compute_fields": list(candidate_fresh_norm_fields), + "candidate_public_raw_total_launch_count": candidate_public_raw_info["runtime_launch_count"], + "baseline_public_raw_assignment_launch_count": baseline_public_raw_fresh_info["assignment_launch_count"], + "baseline_public_raw_norm_launch_count": baseline_public_raw_fresh_info["norm_launch_count"], + "baseline_public_raw_norm_compute_fields": list(baseline_public_raw_fresh_info["norm_compute_fields"]), + "baseline_public_raw_total_launch_count": baseline_public_raw_fresh_info["runtime_launch_count"], + "candidate_precomputed_launch_count": candidate_precomputed_info["prepared_launch_count"], + "candidate_public_raw_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "baseline_public_raw_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "candidate_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "baseline_precomputed_norm_policy": "explicit_precomputed_outside_timing", + "candidate_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_public_raw_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_precomputed_output_policy": "preallocated_outside_timing", + "baseline_precomputed_output_policy": "preallocated_outside_timing", + "candidate_public_raw_shape_first_was_cache_hit": candidate_shape_first_was_cache_hit, + "baseline_public_raw_shape_first_was_cache_hit": baseline_shape_first_was_cache_hit, + "candidate_public_raw_fresh_pointer_cache_hit": bool(candidate_public_raw_info["runtime_cache_hit"]), + "baseline_public_raw_fresh_pointer_cache_hit": bool(baseline_public_raw_fresh_info["runtime_cache_hit"]), + "fresh_pointer_rebind_verified": True, + "candidate_hot_launch_path": candidate_public_raw_info.get("hot_launch_path"), + "candidate_graph_kernel_count": candidate_public_raw_info.get("graph_kernel_count"), + "candidate_graph_capture_error": candidate_public_raw_info.get("graph_capture_error"), + "candidate_public_raw_cold_first_call_activity": candidate_public_raw_shape_info[ + "cold_first_call_activity" + ], + "candidate_prepared_path_proof": prepared_path_proof, + "cold_candidate_public_raw_shape_miss": ( + None if candidate_shape_first_was_cache_hit else _cold_call_fields(candidate_public_raw_shape_cold_call) + ), + "cold_candidate_public_raw_existing_shape_hit": ( + _cold_call_fields(candidate_public_raw_shape_cold_call) if candidate_shape_first_was_cache_hit else None + ), + "cold_candidate_public_raw_fresh_pointer_hit": _cold_call_fields(candidate_public_raw_fresh_cold_call), + "cold_baseline_public_raw_shape_miss": ( + None if baseline_shape_first_was_cache_hit else _cold_call_fields(baseline_public_raw_shape_cold_call) + ), + "cold_baseline_public_raw_existing_shape_hit": ( + _cold_call_fields(baseline_public_raw_shape_cold_call) if baseline_shape_first_was_cache_hit else None + ), + "cold_baseline_public_raw_fresh_pointer_hit": _cold_call_fields(baseline_public_raw_fresh_cold_call), + "cold_candidate_precomputed_prepare_first_pointer": _cold_call_fields(candidate_precomputed_prepare_cold_call), + "cold_candidate_precomputed_prepare_fresh_pointer": _cold_call_fields( + fresh_candidate_precomputed_prepare_cold_call + ), + "cold_candidate_precomputed_first_pointer": _cold_call_fields(candidate_precomputed_first_cold_call), + "cold_candidate_precomputed_fresh_pointer": _cold_call_fields(fresh_candidate_precomputed_first_cold_call), + "cold_baseline_precomputed_first_pointer": _cold_call_fields(baseline_precomputed_first_cold_call), + "cold_baseline_precomputed_fresh_pointer": _cold_call_fields(baseline_precomputed_fresh_cold_call), + } + + if correctness: + ref = _reference_assign(x, centroids, chunk_rows=reference_chunk_rows) + fresh_ref = _reference_assign(fresh_x, fresh_centroids, chunk_rows=reference_chunk_rows) + torch.cuda.synchronize() + baseline_public_raw_correctness = _assignment_correctness( + baseline_public_raw_out, + ref, + x, + centroids, + ) + fresh_baseline_public_raw_correctness = _assignment_correctness( + fresh_baseline_public_raw_out, + fresh_ref, + fresh_x, + fresh_centroids, + ) + candidate_public_raw_correctness = _assignment_correctness( + candidate_public_raw_out, + ref, + x, + centroids, + ) + fresh_candidate_public_raw_correctness = _assignment_correctness( + fresh_candidate_public_raw_out, + fresh_ref, + fresh_x, + fresh_centroids, + ) + candidate_precomputed_correctness = _assignment_correctness( + candidate_precomputed_result, + ref, + x, + centroids, + ) + fresh_candidate_precomputed_correctness = _assignment_correctness( + fresh_candidate_precomputed_result, + fresh_ref, + fresh_x, + fresh_centroids, + ) + baseline_precomputed_correctness = _assignment_correctness( + baseline_precomputed_result, + ref, + x, + centroids, + ) + fresh_baseline_precomputed_correctness = _assignment_correctness( + fresh_baseline_precomputed_result, + fresh_ref, + fresh_x, + fresh_centroids, + ) + result["candidate_public_raw_correctness"] = candidate_public_raw_correctness + result["candidate_public_raw_fresh_pointer_correctness"] = fresh_candidate_public_raw_correctness + result["baseline_public_raw_correctness"] = baseline_public_raw_correctness + result["baseline_public_raw_fresh_pointer_correctness"] = fresh_baseline_public_raw_correctness + result["candidate_precomputed_correctness"] = candidate_precomputed_correctness + result["candidate_precomputed_fresh_pointer_correctness"] = fresh_candidate_precomputed_correctness + result["baseline_precomputed_correctness"] = baseline_precomputed_correctness + result["baseline_precomputed_fresh_pointer_correctness"] = fresh_baseline_precomputed_correctness + correctness_rows = ( + candidate_public_raw_correctness, + fresh_candidate_public_raw_correctness, + baseline_public_raw_correctness, + fresh_baseline_public_raw_correctness, + candidate_precomputed_correctness, + fresh_candidate_precomputed_correctness, + baseline_precomputed_correctness, + fresh_baseline_precomputed_correctness, + ) + result["correct"] = all(bool(item["correct"]) for item in correctness_rows) + + if benchmark: + baseline_public_raw_output_holder = [fresh_baseline_public_raw_out] + candidate_public_raw_output_holder = [fresh_candidate_public_raw_out] + + def held_baseline_public_raw(x_value, centroids_value): + baseline_public_raw_output_holder[0] = baseline_adapter.compute(x_value, centroids_value) + return baseline_public_raw_output_holder[0] + + def held_candidate_public_raw(x_value, centroids_value): + candidate_public_raw_output_holder[0] = runtime.compute(x_value, centroids_value) + return candidate_public_raw_output_holder[0] + + role_factories = { + "candidate_public_raw": lambda: _alternating_call( + lambda: held_candidate_public_raw(x, centroids), + lambda: held_candidate_public_raw(fresh_x, fresh_centroids), + ), + "baseline_public_raw": lambda: _alternating_call( + lambda: held_baseline_public_raw(x, centroids), + lambda: held_baseline_public_raw(fresh_x, fresh_centroids), + ), + "candidate_precomputed": lambda: _alternating_call( + lambda: flash_kmeans_assign_prepared(candidate_prepared), + lambda: flash_kmeans_assign_prepared(fresh_candidate_prepared), + ), + "baseline_precomputed": lambda: _alternating_call( + baseline_precomputed_first_call, + baseline_precomputed_fresh_pointer_call, + ), + } + cold_first_calls = { + "candidate_public_raw": candidate_public_raw_fresh_cold_call, + "baseline_public_raw": baseline_public_raw_fresh_cold_call, + "candidate_precomputed": fresh_candidate_precomputed_first_cold_call, + "baseline_precomputed": baseline_precomputed_fresh_cold_call, + } + timings, paired_session = _bench_interleaved_roles( + label=str(row["label"]), + order_seed=measurement_order_seed, + role_factories=role_factories, + cold_first_calls=cold_first_calls, + ) + candidate_public_raw_timing = timings["candidate_public_raw"] + baseline_public_raw_timing = timings["baseline_public_raw"] + candidate_precomputed_timing = timings["candidate_precomputed"] + baseline_precomputed_timing = timings["baseline_precomputed"] + timing_backends = { + candidate_public_raw_timing.backend, + baseline_public_raw_timing.backend, + candidate_precomputed_timing.backend, + baseline_precomputed_timing.backend, + } + if timing_backends != {"cupti"}: + raise RuntimeError(f"dual-lane benchmark requires CUPTI for every timing block, got {timing_backends}") + result["measurement_order"] = paired_session["first_reportable_roles"] + result["measurement_order_policy"] = "deterministic_interleaved_adjacent_ab_pairs" + result["measurement_order_seed"] = measurement_order_seed + result["paired_cupti_session_id"] = f"{measurement_session_id}:{row['label']}" + result["paired_cupti_session"] = paired_session + result["cold_measurement_order"] = list(cold_measurement_order) + result["cold_measurement_order_policy"] = "deterministic_balanced_per_publication_contract_portfolio" + result["timing_backend"] = "cupti" + result["candidate_public_raw_gpu_span_ms"] = candidate_public_raw_timing.median_gpu_span_ms + result["candidate_public_raw_kernel_sum_ms"] = candidate_public_raw_timing.median_kernel_sum_ms + result["candidate_public_raw_inter_kernel_gap_ms"] = candidate_public_raw_timing.median_inter_kernel_gap_ms + result["candidate_public_raw_submission_ms"] = _required_median( + candidate_public_raw_timing.median_submission_ms, + metric="candidate public-raw submission_ms", + ) + result["candidate_public_raw_host_enqueue_ms"] = _required_median( + candidate_public_raw_timing.median_host_enqueue_ms, + metric="candidate public-raw host_enqueue_ms", + ) + result["candidate_public_raw_synchronized_e2e_ms"] = _required_median( + candidate_public_raw_timing.median_synchronized_e2e_ms, + metric="candidate public-raw synchronized_e2e_ms", + ) + result["candidate_public_raw_timing_backend"] = candidate_public_raw_timing.backend + result["candidate_public_raw_bench_iters"] = len(candidate_public_raw_timing.times_ms) + result["candidate_public_raw_timing_diagnostics"] = _timing_diagnostics( + candidate_public_raw_timing, + primary_metric="synchronized_e2e_ms", + ) + + result["baseline_07cf_adapter_gpu_span_ms"] = baseline_public_raw_timing.median_gpu_span_ms + result["baseline_07cf_adapter_kernel_sum_ms"] = baseline_public_raw_timing.median_kernel_sum_ms + result["baseline_07cf_adapter_inter_kernel_gap_ms"] = baseline_public_raw_timing.median_inter_kernel_gap_ms + result["baseline_07cf_adapter_submission_ms"] = _required_median( + baseline_public_raw_timing.median_submission_ms, + metric="07cf raw-adapter submission_ms", + ) + result["baseline_07cf_adapter_host_enqueue_ms"] = _required_median( + baseline_public_raw_timing.median_host_enqueue_ms, + metric="07cf raw-adapter host_enqueue_ms", + ) + result["baseline_07cf_adapter_synchronized_e2e_ms"] = _required_median( + baseline_public_raw_timing.median_synchronized_e2e_ms, + metric="07cf raw-adapter synchronized_e2e_ms", + ) + result["baseline_07cf_adapter_timing_backend"] = baseline_public_raw_timing.backend + result["baseline_07cf_adapter_bench_iters"] = len(baseline_public_raw_timing.times_ms) + result["baseline_07cf_adapter_timing_diagnostics"] = _timing_diagnostics( + baseline_public_raw_timing, + primary_metric="synchronized_e2e_ms", + ) + + result["candidate_precomputed_gpu_span_ms"] = candidate_precomputed_timing.median_gpu_span_ms + result["candidate_precomputed_kernel_sum_ms"] = candidate_precomputed_timing.median_kernel_sum_ms + result["candidate_precomputed_inter_kernel_gap_ms"] = candidate_precomputed_timing.median_inter_kernel_gap_ms + result["candidate_precomputed_submission_ms"] = _required_median( + candidate_precomputed_timing.median_submission_ms, + metric="candidate prepared submission_ms", + ) + result["candidate_precomputed_host_enqueue_ms"] = _required_median( + candidate_precomputed_timing.median_host_enqueue_ms, + metric="candidate prepared host_enqueue_ms", + ) + result["candidate_precomputed_synchronized_e2e_ms"] = _required_median( + candidate_precomputed_timing.median_synchronized_e2e_ms, + metric="candidate prepared synchronized_e2e_ms", + ) + result["candidate_precomputed_timing_backend"] = candidate_precomputed_timing.backend + result["candidate_precomputed_bench_iters"] = len(candidate_precomputed_timing.times_ms) + result["candidate_precomputed_timing_diagnostics"] = _timing_diagnostics( + candidate_precomputed_timing, + primary_metric="gpu_span_ms", + ) + + result["baseline_07cf_precomputed_gpu_span_ms"] = baseline_precomputed_timing.median_gpu_span_ms + result["baseline_07cf_precomputed_kernel_sum_ms"] = baseline_precomputed_timing.median_kernel_sum_ms + result["baseline_07cf_precomputed_inter_kernel_gap_ms"] = baseline_precomputed_timing.median_inter_kernel_gap_ms + result["baseline_07cf_precomputed_submission_ms"] = _required_median( + baseline_precomputed_timing.median_submission_ms, + metric="07cf prepared submission_ms", + ) + result["baseline_07cf_precomputed_host_enqueue_ms"] = _required_median( + baseline_precomputed_timing.median_host_enqueue_ms, + metric="07cf prepared host_enqueue_ms", + ) + result["baseline_07cf_precomputed_synchronized_e2e_ms"] = _required_median( + baseline_precomputed_timing.median_synchronized_e2e_ms, + metric="07cf prepared synchronized_e2e_ms", + ) + result["baseline_07cf_precomputed_timing_backend"] = baseline_precomputed_timing.backend + result["baseline_07cf_precomputed_bench_iters"] = len(baseline_precomputed_timing.times_ms) + result["baseline_07cf_precomputed_timing_diagnostics"] = _timing_diagnostics( + baseline_precomputed_timing, + primary_metric="gpu_span_ms", + ) + + flops = 2.0 * int(row["B"]) * int(row["N"]) * int(row["K"]) * int(row["D"]) + result["candidate_public_raw_tflops_from_gpu_span"] = ( + flops / candidate_public_raw_timing.median_gpu_span_ms / 1e9 + ) + result.update( + _dual_lane_speedups( + candidate_public_raw_timing=candidate_public_raw_timing, + baseline_public_raw_timing=baseline_public_raw_timing, + candidate_precomputed_timing=candidate_precomputed_timing, + baseline_precomputed_timing=baseline_precomputed_timing, + ) + ) + result["candidate_public_prepared_ratios"] = _public_prepared_ratios( + candidate_public_raw_timing, + candidate_precomputed_timing, + ) + result["baseline_public_prepared_ratios"] = _public_prepared_ratios( + baseline_public_raw_timing, + baseline_precomputed_timing, + ) + candidate_lifecycle = runtime_lifecycle_metrics( + api="flashlib_cake_kmeans.init(...).compute", + measurement_session_id=measurement_session_id, + timing_boundary=CANDIDATE_TIMING_BOUNDARY, + output_policy="default_output_allocated_before_preprocessing_inside_timing", + init=candidate_init_timing, + init_sample_id=measurement_session_id, + first_compute=candidate_public_raw_shape_cold_call, + first_cache_state=("shape_slot_hit" if candidate_shape_first_was_cache_hit else "shape_slot_miss"), + hot_compute=candidate_public_raw_timing, + hot_cache_state="fresh_pointer_shape_slot_hit", + code_cache_state="process_order_dependent", + ) + baseline_lifecycle = runtime_lifecycle_metrics( + api="triton_h200_07cf_raw_adapter_v1.compute", + measurement_session_id=measurement_session_id, + timing_boundary=BASELINE_TIMING_BOUNDARY, + output_policy="default_output_allocated_before_preprocessing_inside_timing", + init=baseline_init_timing, + init_sample_id=measurement_session_id, + first_compute=baseline_public_raw_shape_cold_call, + first_cache_state=("shape_slot_hit" if baseline_shape_first_was_cache_hit else "shape_slot_miss"), + hot_compute=baseline_public_raw_timing, + hot_cache_state="fresh_pointer_shape_slot_hit", + code_cache_state="process_order_dependent", + ) + lifecycle_comparison = compare_runtime_lifecycles(candidate_lifecycle, baseline_lifecycle) + if not math.isclose( + lifecycle_comparison["hot_synchronized_e2e_speedup"], + result["public_raw_e2e_speedup_vs_07cf_adapter"], + rel_tol=1.0e-12, + abs_tol=0.0, + ): + raise RuntimeError("KMeans lifecycle hot E2E speedup disagrees with publication speedup") + result["candidate_runtime_lifecycle"] = candidate_lifecycle + result["baseline_runtime_lifecycle"] = baseline_lifecycle + result["runtime_lifecycle_comparison"] = lifecycle_comparison + + return result + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Correctness and CUPTI benchmark for flashlib_cake_kmeans.flash_kmeans_assign" + ) + parser.add_argument("--shape", action="append", help="Shape label to run. Repeatable.") + parser.add_argument( + "--source", + action="append", + help="Source set to run, for example full95. Repeatable.", + ) + parser.add_argument("--unique", action="store_true", help="Run one row per unique (B,N,D,K,dtype).") + parser.add_argument( + "--include-runtime-coverage", + action="store_true", + help="Include the 104 runtime-only rows; publication defaults to the registry 124.", + ) + parser.add_argument("--limit", type=int, default=None, help="Limit selected rows after filtering.") + parser.add_argument("--arch", default=None, help="NVRTC architecture, e.g. sm_100a.") + parser.add_argument( + "--metadata-only", + action="store_true", + help="Emit available benchmark metadata without CUDA.", + ) + parser.add_argument( + "--no-correctness", + action="store_true", + help="Skip PyTorch reference correctness checks.", + ) + parser.add_argument("--no-benchmark", action="store_true", help="Skip CUPTI timing.") + parser.add_argument( + "--reference-chunk-rows", + type=int, + default=128, + help="Rows per chunk for PyTorch reference.", + ) + parser.add_argument("--json", type=Path, default=None, help="Optional path for JSON output.") + parser.add_argument( + "--order-seed", + default=DEFAULT_MEASUREMENT_ORDER_SEED, + help="Deterministic seed for the interleaved paired CUPTI schedule.", + ) + parser.add_argument("--shard-index", type=int, default=0, help="Zero-based validation shard index.") + parser.add_argument("--shard-count", type=int, default=1, help="Number of validation shards.") + parser.add_argument("--quiet", action="store_true", help="Do not print the full JSON payload.") + args = parser.parse_args() + + if args.limit is not None and args.limit < 0: + parser.error("--limit must be non-negative") + if args.reference_chunk_rows <= 0: + parser.error("--reference-chunk-rows must be positive") + if args.shard_count <= 0 or not 0 <= args.shard_index < args.shard_count: + parser.error("shard index must satisfy 0 <= index < count and count must be positive") + if not args.order_seed: + parser.error("--order-seed must be non-empty") + if args.json is not None: + args.json.unlink(missing_ok=True) + + rows = _selected_rows(args) + measurement_session_id = uuid.uuid4().hex + preprocess_source_digest = _preprocess_source_digest() + if preprocess_source_digest != BASELINE_REGISTRY_PROFILE["shared_preprocess"]["source_sha256"]: + raise RuntimeError("exported preprocessing sources do not match the pinned registry profile") + profile_digest = hashlib.sha256( + json.dumps( + BASELINE_REGISTRY_PROFILE, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + ).encode("utf-8") + ).hexdigest() + if profile_digest != BASELINE_REGISTRY_SHA256: + raise RuntimeError("embedded registry baseline profile digest is stale") + payload: dict[str, Any] = { + "schema": "flash-kmeans-export-dual-lane-v2", + "api": "flashlib_cake_kmeans.init(...).compute", + "semantic_entrypoint": SEMANTIC_ENTRYPOINT, + "publication_metric": "public_raw_e2e_speedup_vs_07cf_adapter", + "parity_metric": "precomputed_gpu_speedup_vs_07cf", + "baseline_name": PUBLIC_RAW_BASELINE_NAME, + "public_raw_baseline_name": PUBLIC_RAW_BASELINE_NAME, + "precomputed_baseline_name": PRECOMPUTED_BASELINE_NAME, + "baseline_commit": BASELINE_COMMIT, + "baseline_entrypoint": BASELINE_ENTRYPOINT, + "benchmark_registry_baseline_key": BASELINE_REGISTRY_KEY, + "benchmark_registry_baseline_sha256": BASELINE_REGISTRY_SHA256, + "benchmark_registry_baseline_profile": BASELINE_REGISTRY_PROFILE, + "registry_candidate_entrypoint": REGISTRY_CANDIDATE_ENTRYPOINT, + "measured_candidate_entrypoint": MEASURED_CANDIDATE_ENTRYPOINT, + "official_timing_boundary": CANDIDATE_TIMING_BOUNDARY, + "candidate_timing_boundary": CANDIDATE_TIMING_BOUNDARY, + "benchmark_registry_candidate_timing_boundary": REGISTRY_CANDIDATE_TIMING_BOUNDARY, + "baseline_timing_boundary": BASELINE_TIMING_BOUNDARY, + "publication_speedup_metric": "public_raw_e2e_speedup_vs_07cf_adapter", + "publication_timing_backend": "cupti", + "publication_speedup_convention": ( + "public_raw_e2e_speedup_vs_07cf_adapter = " + "baseline_07cf_adapter_synchronized_e2e_ms / candidate_public_raw_synchronized_e2e_ms" + ), + "precomputed_parity_speedup_convention": ( + "precomputed_gpu_speedup_vs_07cf = " + "baseline_07cf_precomputed_gpu_span_ms / candidate_precomputed_gpu_span_ms" + ), + "preprocess_impl": PREPROCESS_IMPL, + "preprocess_source_sha256": preprocess_source_digest, + "timing_window_ms": {"warmup_ms": WARMUP_MS, "bench_ms": BENCH_MS}, + "submission_metric": { + "field": "submission_ms", + "host_enqueue_alias_field": "host_enqueue_ms", + "relationship": "same_cpu_submission_bracket", + }, + "artifact": FLASH_KMEANS_EVOLUTION_ARTIFACT, + "evolution_summary": FLASH_KMEANS_EVOLUTION_SUMMARY, + "selected_row_count": len(rows), + "selected_unique_shape_count": len({_shape_key(row) for row in rows}), + "metadata_only": bool(args.metadata_only), + "validation_shard": {"index": args.shard_index, "count": args.shard_count}, + "shapes": rows, + "measurement_session": { + "id": measurement_session_id, + "scope": "per_shape_single_cupti_activity_session_interleaved_paired_roles", + "same_process": True, + "same_cupti_session": ( + True if not args.metadata_only and not args.no_benchmark else None + ), + "interleaved": ( + True if not args.metadata_only and not args.no_benchmark else None + ), + "reportable_timing_collected": bool( + not args.metadata_only and not args.no_benchmark + ), + "randomized_or_interleaved_order": True, + "sequential_full_sweeps": False, + "timing_blocks": "one_interleaved_cupti_activity_session_per_shape", + "adaptive_probe_scope": "separate_nonreportable_cupti_estimation_only", + "alternate_two_pointer_sets": True, + "public_raw_e2e": { + "candidate_api": "flashlib_cake_kmeans.init(...).compute(raw_inputs)", + "baseline_api": "triton_h200_07cf_raw_adapter_v1.compute(raw_inputs)", + "comparison_scope": "complete_raw_input_operators_not_assignment_only", + "candidate_norm_policy": "route_required_internal_fused_bf16_pair_row_norm", + "baseline_norm_policy": "shared_kernel_all_fields_required_by_frozen_07cf", + "candidate_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_scratch_policy": "per_shape_per_stream_cached", + "baseline_scratch_policy": "per_shape_per_stream_cached", + "candidate_runtime_initialized_once": True, + "baseline_runtime_initialized_once": True, + "preprocess_impl": PREPROCESS_IMPL, + "preprocess_source_sha256": preprocess_source_digest, + "shared_preprocess_cold_compile_attributed_to_lane": None, + "fixture_synchronized_before_cold_calls": True, + "assignment_baseline": "frozen_07cf", + }, + "precomputed_kernel_parity": { + "candidate_api": "flash_kmeans_assign_prepared", + "baseline_api": "euclid_assign_triton_h200_07cf", + "candidate_norm_policy": "explicit_precomputed_outside_timing", + "baseline_norm_policy": "explicit_precomputed_outside_timing", + "candidate_output_policy": "preallocated_outside_timing", + "baseline_output_policy": "preallocated_outside_timing", + "both_pointer_sets_prepared_before_timing": True, + }, + "runtime_instances_reused_across_shapes": True, + "resident_multi_shape_cache_benchmarked": False, + "cache_policy": "synchronize_and_clear_after_each_completed_shape", + "baseline_commit": BASELINE_COMMIT, + "order_policy": "deterministic_interleaved_adjacent_ab_pairs", + "cold_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "init_composition": "runtime_init_plus_shared_preprocess_support_each_standalone_lane", + "init_order_policy": "alternate_candidate_baseline_first_by_validation_shard_parity_then_shared_support", + "baseline_has_explicit_init": True, + "order_seed": args.order_seed, + }, + "runtime_lifecycle": { + "schema": "loom-public-runtime-lifecycle-v1", + "candidate_api": "flashlib_cake_kmeans.init(...).compute", + "baseline_api": "triton_h200_07cf_raw_adapter_v1.compute", + "candidate_timing_boundary": CANDIDATE_TIMING_BOUNDARY, + "baseline_timing_boundary": BASELINE_TIMING_BOUNDARY, + "init_scope": ( + "runtime_init_plus_standalone_shared_preprocess_support_once_per_" + "validation_shard_process_device_operator" + ), + "amortization_call_counts": [1, 10, 100, 1000], + "cache_policy": "synchronize_and_clear_after_each_completed_shape", + "cold_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "init_composition": "runtime_init_plus_shared_preprocess_support_each_standalone_lane", + "init_order_policy": "alternate_candidate_baseline_first_by_validation_shard_parity_then_shared_support", + "baseline_has_explicit_init": True, + "resident_multi_shape_cache_benchmarked": False, + "candidate_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "session": None, + }, + } + if args.metadata_only: + payload["results"] = [] + else: + import torch + from flashlib_cake_kmeans import init + + capability = torch.cuda.get_device_capability() + detected_arch = f"sm_{capability[0]}{capability[1]}" + ("" if capability[0] == 8 else "a") + payload["hardware"] = { + "device": torch.cuda.get_device_name(), + "arch": detected_arch, + } + device_index = torch.cuda.current_device() + resolved_arch = args.arch or payload["hardware"]["arch"] + init_measurement_order = ("candidate", "baseline") if args.shard_index % 2 == 0 else ("baseline", "candidate") + if not args.no_benchmark: + require_cupti() + runtime = None + baseline_adapter = None + runtime_init_call = None + baseline_adapter_init_call = None + for lane in init_measurement_order: + if lane == "candidate": + runtime, runtime_init_call = measure_host_call( + lambda: init( + device=device_index, + arch=resolved_arch, + compile="lazy", + ) + ) + else: + baseline_adapter, baseline_adapter_init_call = measure_host_call( + lambda: TritonH20007cfRawAdapter( + device_index=device_index, + arch=resolved_arch, + ) + ) + if runtime is None or baseline_adapter is None: + raise RuntimeError("KMeans init measurement did not construct both standalone lanes") + else: + runtime = init( + device=device_index, + arch=resolved_arch, + compile="lazy", + ) + runtime_init_call = None + baseline_adapter = TritonH20007cfRawAdapter( + device_index=device_index, + arch=resolved_arch, + ) + baseline_adapter_init_call = None + shared_preprocess_cold_compile = ( + _measure_shared_preprocess_cold_compile(rows[0], arch=runtime.arch) + if rows and not args.no_benchmark + else None + ) + candidate_lifecycle_init = ( + _sum_host_timings(runtime_init_call, shared_preprocess_cold_compile) + if runtime_init_call is not None and shared_preprocess_cold_compile is not None + else runtime_init_call + ) + baseline_lifecycle_init = ( + _sum_host_timings(baseline_adapter_init_call, shared_preprocess_cold_compile) + if baseline_adapter_init_call is not None and shared_preprocess_cold_compile is not None + else baseline_adapter_init_call + ) + payload["cold_candidate_runtime_init"] = _cold_call_fields(runtime_init_call) + payload["cold_baseline_07cf_adapter_init"] = _cold_call_fields(baseline_adapter_init_call) + payload["cold_shared_preprocess_compile"] = _cold_call_fields(shared_preprocess_cold_compile) + results = [] + for row in rows: + try: + results.append( + _run_shape( + row, + runtime=runtime, + baseline_adapter=baseline_adapter, + candidate_init_timing=candidate_lifecycle_init, + baseline_init_timing=baseline_lifecycle_init, + arch=args.arch, + correctness=not args.no_correctness, + benchmark=not args.no_benchmark, + reference_chunk_rows=args.reference_chunk_rows, + measurement_session_id=measurement_session_id, + measurement_order_seed=args.order_seed, + ) + ) + finally: + runtime.clear() + baseline_adapter.clear() + payload["results"] = results + payload["runtime_workspace_lifecycle"] = ( + "both_init_once_runtimes_synchronize_and_clear_after_each_completed_shape" + ) + candidate_first_miss_count = sum( + not bool(row["candidate_public_raw_shape_first_was_cache_hit"]) for row in results + ) + baseline_first_miss_count = sum( + not bool(row["baseline_public_raw_shape_first_was_cache_hit"]) for row in results + ) + candidate_fresh_hit_count = sum(bool(row["candidate_public_raw_fresh_pointer_cache_hit"]) for row in results) + baseline_fresh_hit_count = sum(bool(row["baseline_public_raw_fresh_pointer_cache_hit"]) for row in results) + if args.no_benchmark: + payload.pop("runtime_lifecycle", None) + else: + payload["runtime_lifecycle"]["session"] = { + "id": measurement_session_id, + "init_measurement_order": [*init_measurement_order, "shared_preprocess_compile"], + "candidate_init": _cold_call_fields(candidate_lifecycle_init), + "baseline_init": _cold_call_fields(baseline_lifecycle_init), + "init_components": { + "candidate_runtime_init": _cold_call_fields(runtime_init_call), + "baseline_runtime_init": _cold_call_fields(baseline_adapter_init_call), + "shared_preprocess_compile": _cold_call_fields(shared_preprocess_cold_compile), + "attribution": "included_once_in_each_standalone_lane_init_total", + }, + "clear_count": {"candidate": len(results), "baseline": len(results)}, + "first_lookup_miss_count": { + "candidate": candidate_first_miss_count, + "baseline": baseline_first_miss_count, + }, + "fresh_pointer_hit_count": { + "candidate": candidate_fresh_hit_count, + "baseline": baseline_fresh_hit_count, + }, + "final_candidate_cache_info": runtime.cache_info(), + "final_baseline_cache_info": baseline_adapter.cache_info(), + } + + text = json.dumps(payload, indent=2, sort_keys=True) + if args.json is not None: + args.json.parent.mkdir(parents=True, exist_ok=True) + _write_json_atomic(args.json, text) + if not args.quiet: + print(text) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/cake_exports/kmeans/benchmarks/benchmark_exported_kernels.py b/cake_exports/kmeans/benchmarks/benchmark_exported_kernels.py new file mode 100644 index 00000000..4e237f98 --- /dev/null +++ b/cake_exports/kmeans/benchmarks/benchmark_exported_kernels.py @@ -0,0 +1,144 @@ +from __future__ import annotations + +import argparse +import importlib +import json +import statistics +import sys +import time +from pathlib import Path +from typing import Any + + +ROOT = Path(__file__).resolve().parents[1] +SRC = ROOT / "src" +if str(SRC) not in sys.path: + sys.path.insert(0, str(SRC)) + +PACKAGE_NAME = 'flashlib_cake_kmeans' + + +def _selected_names(pkg: Any, requested: list[str] | None) -> list[str]: + if not requested: + return list(pkg.KERNELS) + missing = sorted(set(requested) - set(pkg.KERNELS)) + if missing: + available = ", ".join(sorted(pkg.KERNELS)) + raise SystemExit(f"unknown kernel(s) {missing}. Available: {available}") + return requested + + +def _entry_for_kernel(pkg: Any, name: str, *, metadata_only: bool, arch: str | None, iterations: int): + kernel = pkg.get_kernel(name) + entry: dict[str, Any] = { + "name": name, + "symbol": kernel.spec.symbol, + "launch_mode": kernel.spec.launch_mode, + "threads": kernel.spec.threads, + "shared_mem_bytes": kernel.spec.shared_mem_bytes, + "parameter_count": len(kernel.parameters), + "status": "metadata_only" if metadata_only else "pending", + } + if metadata_only: + return entry + + times_ms: list[float] = [] + try: + runtime = importlib.import_module(f"{PACKAGE_NAME}._runtime") + for _ in range(iterations): + runtime.clear_compilation_cache() + start = time.perf_counter() + kernel.compile(arch=arch) + times_ms.append((time.perf_counter() - start) * 1000.0) + except Exception as exc: # noqa: BLE001 - benchmark report should preserve the failure. + entry["status"] = "failed" + entry["error"] = f"{type(exc).__name__}: {exc}" + return entry + + entry["status"] = "passed" + entry["compile_ms_median"] = statistics.median(times_ms) + entry["compile_ms_min"] = min(times_ms) + entry["compile_ms_max"] = max(times_ms) + entry["iterations"] = iterations + return entry + + +def run_benchmark( + *, + kernels: list[str] | None = None, + arch: str | None = None, + iterations: int = 1, + metadata_only: bool = False, +) -> dict[str, Any]: + pkg = importlib.import_module(PACKAGE_NAME) + selected = _selected_names(pkg, kernels) + entries = [ + _entry_for_kernel(pkg, name, metadata_only=metadata_only, arch=arch, iterations=iterations) + for name in selected + ] + passed = sum(1 for entry in entries if entry["status"] in {"passed", "metadata_only"}) + failed = sum(1 for entry in entries if entry["status"] == "failed") + return { + "benchmark": "exported_kernel_compile", + "package": PACKAGE_NAME, + "arch": arch, + "iterations": iterations, + "metadata_only": metadata_only, + "summary": { + "kernel_count": len(entries), + "passed": passed, + "failed": failed, + "all_passed": failed == 0, + }, + "kernels": entries, + } + + +def _print_summary(payload: dict[str, Any]) -> None: + summary = payload["summary"] + print( + "benchmark={benchmark} package={package} kernels={kernel_count} passed={passed} failed={failed}".format( + benchmark=payload["benchmark"], + package=payload["package"], + kernel_count=summary["kernel_count"], + passed=summary["passed"], + failed=summary["failed"], + ) + ) + for entry in payload["kernels"]: + if entry["status"] == "passed": + print( + "{name}: compile_median={compile_ms_median:.3f} ms status=passed".format(**entry) + ) + else: + detail = f" error={entry['error']}" if "error" in entry else "" + print(f"{entry['name']}: status={entry['status']}{detail}") + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description="Benchmark exported CUDA kernel compile latency.") + parser.add_argument("--kernel", action="append", dest="kernels", help="Kernel name to benchmark.") + parser.add_argument("--arch", help="NVRTC GPU architecture, for example sm_100a.") + parser.add_argument("--iterations", type=int, default=1, help="Compile iterations per kernel.") + parser.add_argument("--metadata-only", action="store_true", help="Do not compile; emit benchmark schema.") + parser.add_argument("--json", type=Path, help="Write benchmark results as JSON.") + args = parser.parse_args(argv) + if args.iterations <= 0: + parser.error("--iterations must be positive") + + payload = run_benchmark( + kernels=args.kernels, + arch=args.arch, + iterations=args.iterations, + metadata_only=args.metadata_only, + ) + if args.json: + args.json.parent.mkdir(parents=True, exist_ok=True) + args.json.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + _print_summary(payload) + return 0 if payload["summary"]["all_passed"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) + diff --git a/cake_exports/kmeans/benchmarks/benchmark_shapes.py b/cake_exports/kmeans/benchmarks/benchmark_shapes.py new file mode 100644 index 00000000..382538d6 --- /dev/null +++ b/cake_exports/kmeans/benchmarks/benchmark_shapes.py @@ -0,0 +1,252 @@ +from __future__ import annotations + +import argparse +import importlib +import json +import sys +from collections.abc import Mapping +from pathlib import Path +from typing import Any + + +ROOT = Path(__file__).resolve().parents[1] +SRC = ROOT / "src" +if str(SRC) not in sys.path: + sys.path.insert(0, str(SRC)) + +PACKAGE_NAME = 'flashlib_cake_kmeans' + + +def _shape_name(shape: Mapping[str, Any], index: int) -> str: + return str(shape.get("name") or f"shape_{index}") + + +def _selected_shapes(workload: Any, requested: list[str] | None) -> list[dict[str, Any]]: + shapes = [dict(shape) for shape in getattr(workload, "SHAPES", ())] + names = [_shape_name(shape, index) for index, shape in enumerate(shapes)] + if len(names) != len(set(names)): + raise ValueError("workload.SHAPES contains duplicate names") + if not requested: + return shapes + missing = sorted(set(requested) - set(names)) + if missing: + raise ValueError(f"unknown shape(s) {missing}. Available: {', '.join(names)}") + requested_set = set(requested) + return [shape for index, shape in enumerate(shapes) if _shape_name(shape, index) in requested_set] + + +def _comparison_payload(value: Any) -> dict[str, Any]: + if isinstance(value, bool): + return {"passed": value} + if not isinstance(value, Mapping) or "passed" not in value: + raise TypeError("case['compare'] must return bool or a mapping containing 'passed'") + return dict(value) + + +def _run_shape( + pkg: Any, + workload: Any, + shape: dict[str, Any], + *, + index: int, + correctness_only: bool, + warmup_iters: int, + bench_iters: int, +) -> dict[str, Any]: + name = _shape_name(shape, index) + entry: dict[str, Any] = {"name": name, "shape": shape, "status": "pending"} + try: + case = workload.make_case(pkg, shape) + if not isinstance(case, Mapping): + raise TypeError("make_case() must return a mapping") + missing = [key for key in ("run", "reference", "compare") if not callable(case.get(key))] + if missing: + raise TypeError(f"benchmark case is missing callable(s): {', '.join(missing)}") + + expected = case["reference"]() + import torch + + cold_first_call = None + if correctness_only: + actual = case["run"]() + torch.cuda.synchronize() + else: + from flashlib_cake_kmeans._benchmark import measure_host_call + + actual, cold_first_call = measure_host_call(case["run"]) + comparison = _comparison_payload(case["compare"](actual, expected)) + entry["correctness"] = comparison + if not comparison["passed"]: + entry["status"] = "incorrect" + return entry + if correctness_only: + entry["status"] = "passed" + return entry + + from flashlib_cake_kmeans._benchmark import bench_gpu_time + + timing = bench_gpu_time( + case["run"], + warmup_iters=warmup_iters, + bench_iters=bench_iters, + cold_l2=True, + cold_first_call=cold_first_call, + ) + entry["timing"] = { + "backend": timing.backend, + "official_gpu_metric": "gpu_span_ms", + "median_ms": timing.median_ms, + "min_ms": timing.min_ms, + "mean_ms": timing.mean_ms, + "iterations": len(timing.times_ms), + "cold_l2": True, + "gpu_span_ms": { + "median": timing.median_gpu_span_ms, + "iterations": timing.times_ms, + }, + "kernel_sum_ms": { + "median": timing.median_kernel_sum_ms, + "iterations": timing.kernel_sum_times_ms, + }, + "active_union_ms": { + "median": timing.median_active_union_ms, + "iterations": timing.active_union_times_ms, + }, + "inter_kernel_gap_ms": { + "median": timing.median_inter_kernel_gap_ms, + "iterations": timing.inter_kernel_gap_times_ms, + }, + "activity_count": { + "median": timing.median_activity_count, + "iterations": timing.activity_counts, + }, + "correlated_launch_activity_count": { + "median": timing.median_launch_activity_count, + "iterations": timing.launch_activity_counts, + }, + "correlated_kernel_activity_count": { + "median": timing.median_kernel_activity_count, + "iterations": timing.kernel_activity_counts, + }, + "host_enqueue_ms": { + "median": timing.median_host_enqueue_ms, + "iterations": timing.host_enqueue_times_ms, + }, + "synchronized_e2e_ms": { + "median": timing.median_synchronized_e2e_ms, + "iterations": timing.synchronized_e2e_times_ms, + }, + "cold_first_call": { + "host_enqueue_ms": timing.cold_first_call_host_enqueue_ms, + "synchronized_e2e_ms": timing.cold_first_call_synchronized_e2e_ms, + }, + } + flops = case.get("flops") + bytes_moved = case.get("bytes") + if flops is not None: + entry["tflops"] = float(flops) / timing.median_ms / 1e9 + if bytes_moved is not None: + entry["gbps"] = float(bytes_moved) / timing.median_ms / 1e6 + if case.get("metrics") is not None: + entry["metrics"] = dict(case["metrics"]) + entry["status"] = "passed" + except Exception as exc: # noqa: BLE001 - preserve per-shape failure in JSON. + entry["status"] = "failed" + entry["error"] = f"{type(exc).__name__}: {exc}" + return entry + + +def run_benchmark( + *, + shapes: list[str] | None = None, + metadata_only: bool = False, + correctness_only: bool = False, + warmup_iters: int = 5, + bench_iters: int = 20, +) -> dict[str, Any]: + pkg = importlib.import_module(PACKAGE_NAME) + workload = importlib.import_module("workload") + selected = _selected_shapes(workload, shapes) + configured = bool(getattr(workload, "CONFIGURED", bool(selected))) + entries = [ + {"name": _shape_name(shape, index), "shape": shape, "status": "metadata_only"} + for index, shape in enumerate(selected) + ] + if not metadata_only: + if not configured or not selected: + raise RuntimeError( + "benchmarks/workload.py is not configured; provide --benchmark-adapter during export " + "or implement SHAPES and make_case()" + ) + if not correctness_only: + # CUPTI must be imported before workload adapters import + # torch, which may otherwise load an incompatible system + # CUPTI soname first. + benchmark_runtime = importlib.import_module(f"{PACKAGE_NAME}._benchmark") + benchmark_runtime.require_cupti() + entries = [ + _run_shape( + pkg, + workload, + shape, + index=index, + correctness_only=correctness_only, + warmup_iters=warmup_iters, + bench_iters=bench_iters, + ) + for index, shape in enumerate(selected) + ] + passed = sum(1 for entry in entries if entry["status"] in {"passed", "metadata_only"}) + failed = len(entries) - passed + return { + "benchmark": "exported_kernel_shapes", + "package": PACKAGE_NAME, + "metadata_only": metadata_only, + "correctness_only": correctness_only, + "adapter_configured": configured, + "timing_backend_requested": "cupti", + "summary": { + "shape_count": len(entries), + "passed": passed, + "failed": failed, + "all_passed": failed == 0, + }, + "shapes": entries, + } + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description="Validate and benchmark exported kernels by shape.") + parser.add_argument("--shape", action="append", dest="shapes", help="Shape name to run.") + parser.add_argument("--metadata-only", action="store_true", help="Emit adapter metadata without CUDA.") + parser.add_argument("--correctness-only", action="store_true", help="Validate without timing.") + parser.add_argument("--warmup-iters", type=int, default=5) + parser.add_argument("--bench-iters", type=int, default=20) + parser.add_argument("--json", type=Path, help="Write benchmark results as JSON.") + args = parser.parse_args(argv) + if args.warmup_iters < 0 or args.bench_iters <= 0: + parser.error("--warmup-iters must be non-negative and --bench-iters must be positive") + try: + payload = run_benchmark( + shapes=args.shapes, + metadata_only=args.metadata_only, + correctness_only=args.correctness_only, + warmup_iters=args.warmup_iters, + bench_iters=args.bench_iters, + ) + except Exception as exc: # noqa: BLE001 - CLI reports configuration failures cleanly. + parser.error(str(exc)) + if args.json: + args.json.parent.mkdir(parents=True, exist_ok=True) + args.json.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + summary = payload["summary"] + print( + f"benchmark={payload['benchmark']} shapes={summary['shape_count']} " + f"passed={summary['passed']} failed={summary['failed']}" + ) + return 0 if summary["all_passed"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) + diff --git a/cake_exports/kmeans/benchmarks/expected_routes.json b/cake_exports/kmeans/benchmarks/expected_routes.json new file mode 100644 index 00000000..9c74f027 --- /dev/null +++ b/cake_exports/kmeans/benchmarks/expected_routes.json @@ -0,0 +1,622 @@ +[ + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_1d49_d112_highk_tail_b1_n1408_k4096_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_1d49_d128_forced_fallback_b5_n7296_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_1d49_d224_highn_overlap_b5_n5632_k768_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_1d49_d288_low_n_heldout_b3_n1152_k2048_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d352_exactd_splitk_c95c_v2:launch_for_eval", + "selected_route": "d352_exactd_splitk_c95c_v2", + "shape": "adjacent_1d49_d352_request_highk_b2_n1024_k8192_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_1d49_d416_random_b4_n3840_k512_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_1d49_d480_smallk_boundary_b2_n2816_k256_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_1d49_d48_lowk_boundary_b4_n3968_k256_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1:launch_for_eval", + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "shape": "adjacent_3328_d112_highk_low_n_b2_n512_k8192_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "adjacent_3328_d128_exact_guard_k_neighbor_b8_n8192_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_3328_d224_tail_div_b3_n3840_k512_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_3328_d288_guard_overlap_b4_n8192_k256_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_3328_d352_random_legal_b3_n2048_k768_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_3328_d416_highk_low_n_b1_n384_k8192_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_3328_d480_min_boundary_b1_n128_k256_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_3328_d48_small_boundary_b1_n256_k256_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_5600_d112_odd_tile_tail_b5_n2176_k512_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_5600_d128_fallback_neighbor_b8_n8320_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_5600_d224_k512_overlap_b3_n3072_k512_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_5600_d288_mid_highk_b2_n1024_k2048_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d352_exactd_splitk_c95c_v2:launch_for_eval", + "selected_route": "d352_exactd_splitk_c95c_v2", + "shape": "adjacent_5600_d352_splitk_edge_b2_n768_k4096_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_5600_d416_smallk_boundary_b4_n2048_k256_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d480_splitk_k1024_eac2_v1:launch_for_eval", + "selected_route": "d480_splitk_k1024_eac2_v1", + "shape": "adjacent_5600_d480_random_b1_n1536_k1024_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_5600_d48_low_n_boundary_b3_n1536_k256_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_68cf_d112_tail_b5_n2944_k512_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_68cf_d128_forced_fallback_b7_n6016_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_68cf_d224_overlap_b3_n5120_k1024_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_68cf_d288_boundary_b2_n1920_k512_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_68cf_d352_tail_b3_n2816_k768_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_68cf_d416_overlap_b2_n2304_k1024_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_68cf_d480_boundary_b4_n1664_k512_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_68cf_d48_boundary_b4_n2304_k512_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1:launch_for_eval", + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "shape": "adjacent_8f09_d112_small_highk_b1_n384_k4096_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_8f09_d128_fallback_neighbor_b4_n4480_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4:launch_for_eval", + "selected_route": "d224_tmem_abi_repair_d17c_v4", + "shape": "adjacent_8f09_d224_low_n_boundary_b4_n1536_k256_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_8f09_d288_k768_overlap_b1_n2560_k768_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_8f09_d352_smallk_large_n_b4_n4096_k256_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_8f09_d416_midk_tail_b3_n3456_k768_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d480_splitk_k1024_eac2_v1:launch_for_eval", + "selected_route": "d480_splitk_k1024_eac2_v1", + "shape": "adjacent_8f09_d480_highk_low_n_b2_n640_k4096_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_8f09_d48_medium_tail_b2_n1792_k512_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_9ca0_d112_tail_div_b3_n3840_k768_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_9ca0_d128_exact_guard_neighbor_b8_n8064_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_9ca0_d224_guard_overlap_b2_n4096_k256_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_9ca0_d288_highk_low_n_b1_n384_k4096_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d352_exactd_splitk_c95c_v2:launch_for_eval", + "selected_route": "d352_exactd_splitk_c95c_v2", + "shape": "adjacent_9ca0_d352_splitk_boundary_b1_n512_k8192_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_9ca0_d416_splitk_min_b1_n1024_k4096_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_9ca0_d480_splitd_large_n_b2_n4096_k512_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_9ca0_d48_guard_boundary_b6_n12288_k512_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_a2f8_d112_lowk_tail_b4_n3456_k256_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_a2f8_d128_odd_fallback_b6_n8576_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_a2f8_d224_highn_overlap_b2_n6144_k768_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_a2f8_d288_splitk_probe_b1_n640_k4096_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_a2f8_d352_smallk_boundary_b4_n1024_k256_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_a2f8_d416_random_b2_n2560_k768_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d480_splitk_k1024_eac2_v1:launch_for_eval", + "selected_route": "d480_splitk_k1024_eac2_v1", + "shape": "adjacent_a2f8_d480_highk_request_b1_n896_k4096_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_a2f8_d48_k1024_boundary_b2_n768_k1024_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_c44f_d112_odd_tail_b2_n3200_k768_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_c44f_d128_forced_fallback_b6_n6272_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_c44f_d224_overlap_b4_n4480_k512_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_c44f_d288_highk_heldout_b2_n768_k4096_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_c44f_d352_random_b1_n3328_k768_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_c44f_d416_highk_request_b3_n512_k8192_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_c44f_d480_boundary_b5_n2048_k512_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_c44f_d48_midk_boundary_b3_n2688_k768_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1:launch_for_eval", + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "shape": "adjacent_d9d5_d112_highk_request_b3_n768_k8192_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_d9d5_d128_guard_miss_fallback_b5_n6016_k1024_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_d9d5_d224_random_midk_b2_n2944_k1280_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_d9d5_d288_heldout_low_n_b1_n896_k4096_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_d9d5_d352_random_b5_n2304_k768_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_d9d5_d416_highk_request_b2_n640_k8192_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d480_splitk_k1024_eac2_v1:launch_for_eval", + "selected_route": "d480_splitk_k1024_eac2_v1", + "shape": "adjacent_d9d5_d480_heldout_highk_b3_n1024_k4096_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_d9d5_d48_highk_heldout_b2_n640_k4096_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_ef00_d112_odd_tail_b4_n3712_k1024_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_ef00_d128_forced_fallback_b4_n7552_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_ef00_d224_midn_overlap_b2_n5376_k512_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_ef00_d288_highk_heldout_b1_n1664_k4096_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d352_exactd_splitk_c95c_v2:launch_for_eval", + "selected_route": "d352_exactd_splitk_c95c_v2", + "shape": "adjacent_ef00_d352_request_low_n_b3_n896_k8192_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_ef00_d416_random_midk_b1_n2176_k768_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_ef00_d480_lowk_boundary_b3_n3200_k256_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_ef00_d48_midn_boundary_b3_n3456_k512_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1:launch_for_eval", + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "shape": "post_d895_d112_b1_n256_k256_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1:launch_for_eval", + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "shape": "post_d895_d112_b1_n512_k8192_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1:launch_for_eval", + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "shape": "post_d895_d112_b2_n2048_k512_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d112_b4_n8192_k1024_d112" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "post_d895_d128_fallback_b3_n1920_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "post_d895_d128_fallback_b5_n2176_k512_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "aligned_weave_v10_fallback", + "shape": "post_d895_d128_fallback_b7_n2432_k1024_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval", + "selected_route": "paired_large_v15", + "shape": "post_d895_d128_paired_b2_n262144_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval", + "selected_route": "d128_even_near_floor_v10_repair", + "shape": "post_d895_d128_paired_b8_n8192_k256_d128" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1:launch_for_eval", + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d144_b1_n256_k256_d144" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1:launch_for_eval", + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d144_b1_n512_k8192_d144" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1:launch_for_eval", + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d144_b2_n2048_k1024_d144" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1:launch_for_eval", + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d144_b4_n8192_k1024_d144" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_microdim_hybrid_9c0d_v1:launch_for_eval", + "selected_route": "microdim_hybrid_9c0d_v1", + "shape": "post_d895_d16_b4_n32768_k1024_d16" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_microdim_pipeline4_08f9_v4:launch_for_eval", + "selected_route": "microdim_pipeline4_08f9_v4", + "shape": "post_d895_d16_b8_n65536_k512_d16" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1:launch_for_eval", + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d176_b1_n256_k256_d176" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1:launch_for_eval", + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d176_b1_n512_k8192_d176" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1:launch_for_eval", + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d176_b2_n2048_k1024_d176" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1:launch_for_eval", + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d176_b4_n8192_k1024_d176" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d224_b1_n256_k256_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d224_b1_n512_k8192_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d224_b2_n2048_k1024_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d224_b4_n8192_k1024_d224" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "post_d895_d288_b1_n256_k256_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "post_d895_d288_b1_n512_k8192_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "post_d895_d288_b2_n2048_k1024_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval", + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "post_d895_d288_b4_n8192_k1024_d288" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_microdim_hybrid_9c0d_v1:launch_for_eval", + "selected_route": "microdim_hybrid_9c0d_v1", + "shape": "post_d895_d32_b4_n32768_k1024_d32" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_microdim_pipeline4_08f9_v4:launch_for_eval", + "selected_route": "microdim_pipeline4_08f9_v4", + "shape": "post_d895_d32_b8_n65536_k512_d32" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d352_b1_n256_k256_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d352_exactd_splitk_c95c_v2:launch_for_eval", + "selected_route": "d352_exactd_splitk_c95c_v2", + "shape": "post_d895_d352_b1_n512_k8192_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d352_exactd_splitk_c95c_v2:launch_for_eval", + "selected_route": "d352_exactd_splitk_c95c_v2", + "shape": "post_d895_d352_b2_n2048_k1024_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d352_b4_n8192_k1024_d352" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d416_b1_n256_k256_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d416_b1_n512_k8192_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d416_b2_n2048_k1024_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d416_b4_n8192_k1024_d416" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d480_b1_n256_k256_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d480_splitk_k1024_eac2_v1:launch_for_eval", + "selected_route": "d480_splitk_k1024_eac2_v1", + "shape": "post_d895_d480_b1_n512_k8192_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_d480_splitk_k1024_eac2_v1:launch_for_eval", + "selected_route": "d480_splitk_k1024_eac2_v1", + "shape": "post_d895_d480_b2_n2048_k1024_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d480_b4_n8192_k1024_d480" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d48_b1_n512_k8192_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d48_b2_n2048_k512_d48" + }, + { + "launch_entrypoint": "loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval", + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d48_b4_n8192_k1024_d48" + } +] diff --git a/cake_exports/kmeans/benchmarks/flash_kmeans_shapes.py b/cake_exports/kmeans/benchmarks/flash_kmeans_shapes.py new file mode 100644 index 00000000..a66560b3 --- /dev/null +++ b/cake_exports/kmeans/benchmarks/flash_kmeans_shapes.py @@ -0,0 +1,3436 @@ +from __future__ import annotations + +FLASH_KMEANS_EVOLUTION_ARTIFACT = ( + "artifacts/generalize_auto_tuning/flash_kmeans_assign_20260620/dispatcher_perf_all_shapes_latest.md" +) + +FLASH_KMEANS_EVOLUTION_SUMMARY = { + "mean_evolution_speedup": 19.980689035087718, + "mean_evolution_tflops": 25.401755701754386, + "min_evolution_speedup": 1.3307, + "row_count": 228, + "unique_shape_count": 225, +} + +FLASH_KMEANS_SHAPES = [ + { + "source": "expanded9", + "label": "d895_expanded_heldout_neighborhood_d80_b2_n2176_k1024_d80", + "B": 2, + "N": 2176, + "D": 80, + "K": 1024, + "dtype": "bfloat16", + "route": "lowdim_e50c_v1", + "evolution_kernel_ms": 0.176928, + "evolution_flashlib_ms": 0.402944, + "evolution_tflops": 4.0301, + "evolution_flashlib_equiv_tflops": 1.7696, + "evolution_speedup": 2.2774, + }, + { + "source": "expanded9", + "label": "d895_expanded_guard_miss_fallback_d128_b1_n1408_k512_d128", + "B": 1, + "N": 1408, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.144544, + "evolution_flashlib_ms": 0.348175, + "evolution_tflops": 1.2768, + "evolution_flashlib_equiv_tflops": 0.53, + "evolution_speedup": 2.4088, + }, + { + "source": "expanded9", + "label": "d895_expanded_forced_fallback_d128_b1_n1664_k256_d128", + "B": 1, + "N": 1664, + "D": 128, + "K": 256, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.142032, + "evolution_flashlib_ms": 0.277904, + "evolution_tflops": 0.7678, + "evolution_flashlib_equiv_tflops": 0.3924, + "evolution_speedup": 1.9566, + }, + { + "source": "expanded9", + "label": "d895_expanded_guard_boundary_d144_b1_n128_k256_d144", + "B": 1, + "N": 128, + "D": 144, + "K": 256, + "dtype": "bfloat16", + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_kernel_ms": 0.1664, + "evolution_flashlib_ms": 0.459647, + "evolution_tflops": 0.0567, + "evolution_flashlib_equiv_tflops": 0.0205, + "evolution_speedup": 2.7623, + }, + { + "source": "expanded9", + "label": "d895_expanded_tail_divisibility_d176_b1_n1152_k512_d176", + "B": 1, + "N": 1152, + "D": 176, + "K": 512, + "dtype": "bfloat16", + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_kernel_ms": 0.174944, + "evolution_flashlib_ms": 0.395232, + "evolution_tflops": 1.1868, + "evolution_flashlib_equiv_tflops": 0.5253, + "evolution_speedup": 2.2592, + }, + { + "source": "expanded9", + "label": "d895_expanded_guard_overlap_d192_b1_n1024_k768_d192", + "B": 1, + "N": 1024, + "D": 192, + "K": 768, + "dtype": "bfloat16", + "route": "d192_paired_repeated_mma_v1", + "evolution_kernel_ms": 0.17168, + "evolution_flashlib_ms": 0.45632, + "evolution_tflops": 1.759, + "evolution_flashlib_equiv_tflops": 0.6618, + "evolution_speedup": 2.658, + }, + { + "source": "expanded9", + "label": "d895_expanded_random_legal_d320_b1_n1280_k512_d320", + "B": 1, + "N": 1280, + "D": 320, + "K": 512, + "dtype": "bfloat16", + "route": "highd_splitd_single_tile_6fcf_v1", + "evolution_kernel_ms": 0.152496, + "evolution_flashlib_ms": 0.507759, + "evolution_tflops": 2.7504, + "evolution_flashlib_equiv_tflops": 0.826, + "evolution_speedup": 3.3297, + }, + { + "source": "expanded9", + "label": "d895_expanded_request_specific_d448_b1_n1280_k1024_d448", + "B": 1, + "N": 1280, + "D": 448, + "K": 1024, + "dtype": "bfloat16", + "route": "highd_splitd_single_tile_6fcf_v1", + "evolution_kernel_ms": 0.173743, + "evolution_flashlib_ms": 0.674512, + "evolution_tflops": 6.7594, + "evolution_flashlib_equiv_tflops": 1.7411, + "evolution_speedup": 3.8822, + }, + { + "source": "expanded9", + "label": "d895_expanded_request_specific_d512_splitk_b1_n512_k8192_d512", + "B": 1, + "N": 512, + "D": 512, + "K": 8192, + "dtype": "bfloat16", + "route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "evolution_kernel_ms": 0.257392, + "evolution_flashlib_ms": 31.230383, + "evolution_tflops": 16.6865, + "evolution_flashlib_equiv_tflops": 0.1375, + "evolution_speedup": 121.3339, + }, + { + "source": "full95", + "label": "d16_paired_b2_n4096_k1024_d16", + "B": 2, + "N": 4096, + "D": 16, + "K": 1024, + "dtype": "bfloat16", + "route": "microdim_hybrid_9c0d_v1", + "evolution_kernel_ms": 0.175568, + "evolution_flashlib_ms": 0.314751, + "evolution_tflops": 1.529, + "evolution_flashlib_equiv_tflops": 0.8529, + "evolution_speedup": 1.7928, + }, + { + "source": "full95", + "label": "d16_fallback_b3_n2432_k512_d16", + "B": 3, + "N": 2432, + "D": 16, + "K": 512, + "dtype": "bfloat16", + "route": "microdim_hybrid_9c0d_v1", + "evolution_kernel_ms": 0.127456, + "evolution_flashlib_ms": 0.291712, + "evolution_tflops": 0.9379, + "evolution_flashlib_equiv_tflops": 0.4098, + "evolution_speedup": 2.2887, + }, + { + "source": "full95", + "label": "d16_small_b4_n1024_k512_d16", + "B": 4, + "N": 1024, + "D": 16, + "K": 512, + "dtype": "bfloat16", + "route": "microdim_hybrid_9c0d_v1", + "evolution_kernel_ms": 0.141568, + "evolution_flashlib_ms": 0.365599, + "evolution_tflops": 0.474, + "evolution_flashlib_equiv_tflops": 0.1836, + "evolution_speedup": 2.5825, + }, + { + "source": "full95", + "label": "d16_large_b8_n32768_k512_d16", + "B": 8, + "N": 32768, + "D": 16, + "K": 512, + "dtype": "bfloat16", + "route": "microdim_hybrid_9c0d_v1", + "evolution_kernel_ms": 0.234416, + "evolution_flashlib_ms": 0.3296, + "evolution_tflops": 18.322, + "evolution_flashlib_equiv_tflops": 13.0308, + "evolution_speedup": 1.406, + }, + { + "source": "full95", + "label": "d32_hugek_b1_n512_k8192_d32", + "B": 1, + "N": 512, + "D": 32, + "K": 8192, + "dtype": "bfloat16", + "route": "microdim_hybrid_9c0d_v1", + "evolution_kernel_ms": 0.224064, + "evolution_flashlib_ms": 1.396063, + "evolution_tflops": 1.198, + "evolution_flashlib_equiv_tflops": 0.1923, + "evolution_speedup": 6.2306, + }, + { + "source": "full95", + "label": "d32_paired_b2_n4096_k1024_d32", + "B": 2, + "N": 4096, + "D": 32, + "K": 1024, + "dtype": "bfloat16", + "route": "microdim_hybrid_9c0d_v1", + "evolution_kernel_ms": 0.176672, + "evolution_flashlib_ms": 0.317696, + "evolution_tflops": 3.0388, + "evolution_flashlib_equiv_tflops": 1.6899, + "evolution_speedup": 1.7982, + }, + { + "source": "full95", + "label": "d32_fallback_b3_n2432_k512_d32", + "B": 3, + "N": 2432, + "D": 32, + "K": 512, + "dtype": "bfloat16", + "route": "microdim_hybrid_9c0d_v1", + "evolution_kernel_ms": 0.133824, + "evolution_flashlib_ms": 0.28208, + "evolution_tflops": 1.7865, + "evolution_flashlib_equiv_tflops": 0.8475, + "evolution_speedup": 2.1078, + }, + { + "source": "full95", + "label": "d32_small_b4_n1024_k512_d32", + "B": 4, + "N": 1024, + "D": 32, + "K": 512, + "dtype": "bfloat16", + "route": "microdim_hybrid_9c0d_v1", + "evolution_kernel_ms": 0.139583, + "evolution_flashlib_ms": 0.361167, + "evolution_tflops": 0.9616, + "evolution_flashlib_equiv_tflops": 0.3716, + "evolution_speedup": 2.5875, + }, + { + "source": "full95", + "label": "d64_paired_b2_n4096_k1024_d64", + "B": 2, + "N": 4096, + "D": 64, + "K": 1024, + "dtype": "bfloat16", + "route": "d64_direct_single64_1p2gap_9f2a_v1", + "evolution_kernel_ms": 0.070272, + "evolution_flashlib_ms": 0.320352, + "evolution_tflops": 15.2798, + "evolution_flashlib_equiv_tflops": 3.3518, + "evolution_speedup": 4.5587, + }, + { + "source": "full95", + "label": "d64_fallback_b3_n2432_k512_d64", + "B": 3, + "N": 2432, + "D": 64, + "K": 512, + "dtype": "bfloat16", + "route": "d64_direct_single64_1p2gap_9f2a_v1", + "evolution_kernel_ms": 0.067552, + "evolution_flashlib_ms": 0.2856, + "evolution_tflops": 7.0783, + "evolution_flashlib_equiv_tflops": 1.6742, + "evolution_speedup": 4.2279, + }, + { + "source": "full95", + "label": "d64_small_b4_n1024_k512_d64", + "B": 4, + "N": 1024, + "D": 64, + "K": 512, + "dtype": "bfloat16", + "route": "d64_direct_single64_1p2gap_9f2a_v1", + "evolution_kernel_ms": 0.06736, + "evolution_flashlib_ms": 0.398223, + "evolution_tflops": 3.9851, + "evolution_flashlib_equiv_tflops": 0.6741, + "evolution_speedup": 5.9119, + }, + { + "source": "full95", + "label": "d64_large_b8_n32768_k512_d64", + "B": 8, + "N": 32768, + "D": 64, + "K": 512, + "dtype": "bfloat16", + "route": "d64_direct_single64_1p2gap_9f2a_v1", + "evolution_kernel_ms": 0.128672, + "evolution_flashlib_ms": 0.493615, + "evolution_tflops": 133.5168, + "evolution_flashlib_equiv_tflops": 34.8042, + "evolution_speedup": 3.8362, + }, + { + "source": "full95", + "label": "d80_paired_b1_n2048_k1024_d80", + "B": 1, + "N": 2048, + "D": 80, + "K": 1024, + "dtype": "bfloat16", + "route": "lowdim_e50c_v1", + "evolution_kernel_ms": 0.176896, + "evolution_flashlib_ms": 0.384112, + "evolution_tflops": 1.8968, + "evolution_flashlib_equiv_tflops": 0.8736, + "evolution_speedup": 2.1714, + }, + { + "source": "full95", + "label": "d80_small_b2_n1024_k512_d80", + "B": 2, + "N": 1024, + "D": 80, + "K": 512, + "dtype": "bfloat16", + "route": "lowdim_e50c_v1", + "evolution_kernel_ms": 0.173552, + "evolution_flashlib_ms": 0.378656, + "evolution_tflops": 0.9667, + "evolution_flashlib_equiv_tflops": 0.4431, + "evolution_speedup": 2.1818, + }, + { + "source": "full95", + "label": "d96_paired_b1_n1536_k1024_d96", + "B": 1, + "N": 1536, + "D": 96, + "K": 1024, + "dtype": "bfloat16", + "route": "lowdim_e50c_v1", + "evolution_kernel_ms": 0.176032, + "evolution_flashlib_ms": 0.427679, + "evolution_tflops": 1.7155, + "evolution_flashlib_equiv_tflops": 0.7061, + "evolution_speedup": 2.4296, + }, + { + "source": "full95", + "label": "d96_small_b2_n896_k512_d96", + "B": 2, + "N": 896, + "D": 96, + "K": 512, + "dtype": "bfloat16", + "route": "lowdim_e50c_v1", + "evolution_kernel_ms": 0.173504, + "evolution_flashlib_ms": 0.394576, + "evolution_tflops": 1.0153, + "evolution_flashlib_equiv_tflops": 0.4465, + "evolution_speedup": 2.2742, + }, + { + "source": "full95", + "label": "d96_random_b3_n6144_k1280_d96", + "B": 3, + "N": 6144, + "D": 96, + "K": 1280, + "dtype": "bfloat16", + "route": "lowdim_e50c_v1", + "evolution_kernel_ms": 0.177823, + "evolution_flashlib_ms": 0.426416, + "evolution_tflops": 25.4739, + "evolution_flashlib_equiv_tflops": 10.6231, + "evolution_speedup": 2.398, + }, + { + "source": "full95", + "label": "d96_fallback_b5_n1664_k256_d96", + "B": 5, + "N": 1664, + "D": 96, + "K": 256, + "dtype": "bfloat16", + "route": "lowdim_e50c_v1", + "evolution_kernel_ms": 0.171136, + "evolution_flashlib_ms": 0.304464, + "evolution_tflops": 2.3896, + "evolution_flashlib_equiv_tflops": 1.3432, + "evolution_speedup": 1.7791, + }, + { + "source": "full95", + "label": "boundary_b1_n128_k256_d128", + "B": 1, + "N": 128, + "D": 128, + "K": 256, + "dtype": "bfloat16", + "route": "small_grid_single_tile_v10", + "evolution_kernel_ms": 0.143039, + "evolution_flashlib_ms": 0.384415, + "evolution_tflops": 0.0586, + "evolution_flashlib_equiv_tflops": 0.0218, + "evolution_speedup": 2.6875, + }, + { + "source": "full95", + "label": "fallback_tiny_hugek_b1_n128_k4096_d128", + "B": 1, + "N": 128, + "D": 128, + "K": 4096, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.171472, + "evolution_flashlib_ms": 1.205311, + "evolution_tflops": 0.7827, + "evolution_flashlib_equiv_tflops": 0.1114, + "evolution_speedup": 7.0292, + }, + { + "source": "full95", + "label": "small_b1_n256_k256_d128", + "B": 1, + "N": 256, + "D": 128, + "K": 256, + "dtype": "bfloat16", + "route": "small_grid_single_tile_v10", + "evolution_kernel_ms": 0.145312, + "evolution_flashlib_ms": 0.435263, + "evolution_tflops": 0.1155, + "evolution_flashlib_equiv_tflops": 0.0385, + "evolution_speedup": 2.9954, + }, + { + "source": "full95", + "label": "paired_min_even_b1_n256_k768_d128", + "B": 1, + "N": 256, + "D": 128, + "K": 768, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.167616, + "evolution_flashlib_ms": 0.599455, + "evolution_tflops": 0.3003, + "evolution_flashlib_equiv_tflops": 0.084, + "evolution_speedup": 3.5764, + }, + { + "source": "full95", + "label": "paired_low_n_hugek_b1_n256_k4096_d128", + "B": 1, + "N": 256, + "D": 128, + "K": 4096, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.209535, + "evolution_flashlib_ms": 1.432831, + "evolution_tflops": 1.2811, + "evolution_flashlib_equiv_tflops": 0.1873, + "evolution_speedup": 6.8381, + }, + { + "source": "full95", + "label": "small_b1_n896_k512_d128", + "B": 1, + "N": 896, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "small_grid_single_tile_v10", + "evolution_kernel_ms": 0.144463, + "evolution_flashlib_ms": 0.378655, + "evolution_tflops": 0.8129, + "evolution_flashlib_equiv_tflops": 0.3102, + "evolution_speedup": 2.6211, + }, + { + "source": "full95", + "label": "fallback_odd_kover_b1_n896_k768_d128", + "B": 1, + "N": 896, + "D": 128, + "K": 768, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.145856, + "evolution_flashlib_ms": 0.446111, + "evolution_tflops": 1.2078, + "evolution_flashlib_equiv_tflops": 0.3949, + "evolution_speedup": 3.0586, + }, + { + "source": "full95", + "label": "fallback_hugek_b1_n896_k16384_d128", + "B": 1, + "N": 896, + "D": 128, + "K": 16384, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.25696, + "evolution_flashlib_ms": 2.21235, + "evolution_tflops": 14.6252, + "evolution_flashlib_equiv_tflops": 1.6987, + "evolution_speedup": 8.6097, + }, + { + "source": "full95", + "label": "boundary_b1_n1024_k768_d128", + "B": 1, + "N": 1024, + "D": 128, + "K": 768, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.169952, + "evolution_flashlib_ms": 0.466272, + "evolution_tflops": 1.1846, + "evolution_flashlib_equiv_tflops": 0.4318, + "evolution_speedup": 2.7435, + }, + { + "source": "full95", + "label": "paired_kover_b1_n1024_k2048_d128", + "B": 1, + "N": 1024, + "D": 128, + "K": 2048, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.1848, + "evolution_flashlib_ms": 0.603104, + "evolution_tflops": 2.9051, + "evolution_flashlib_equiv_tflops": 0.8902, + "evolution_speedup": 3.2635, + }, + { + "source": "full95", + "label": "forced_fallback_b1_n1152_k256_d128", + "B": 1, + "N": 1152, + "D": 128, + "K": 256, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.143648, + "evolution_flashlib_ms": 0.313135, + "evolution_tflops": 0.5256, + "evolution_flashlib_equiv_tflops": 0.2411, + "evolution_speedup": 2.1799, + }, + { + "source": "full95", + "label": "fallback_first_after_small_b1_n1152_k512_d128", + "B": 1, + "N": 1152, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.144832, + "evolution_flashlib_ms": 0.38096, + "evolution_tflops": 1.0426, + "evolution_flashlib_equiv_tflops": 0.3964, + "evolution_speedup": 2.6304, + }, + { + "source": "full95", + "label": "boundary_b1_n1280_k256_d128", + "B": 1, + "N": 1280, + "D": 128, + "K": 256, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.163984, + "evolution_flashlib_ms": 0.308176, + "evolution_tflops": 0.5116, + "evolution_flashlib_equiv_tflops": 0.2722, + "evolution_speedup": 1.8793, + }, + { + "source": "full95", + "label": "paired_even_b1_n1536_k256_d128", + "B": 1, + "N": 1536, + "D": 128, + "K": 256, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.160624, + "evolution_flashlib_ms": 0.276479, + "evolution_tflops": 0.6267, + "evolution_flashlib_equiv_tflops": 0.3641, + "evolution_speedup": 1.7213, + }, + { + "source": "full95", + "label": "paired_largek_b1_n2048_k4096_d128", + "B": 1, + "N": 2048, + "D": 128, + "K": 4096, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.209024, + "evolution_flashlib_ms": 0.729904, + "evolution_tflops": 10.2739, + "evolution_flashlib_equiv_tflops": 2.9421, + "evolution_speedup": 3.492, + }, + { + "source": "full95", + "label": "fallback_odd_b1_n2176_k4096_d128", + "B": 1, + "N": 2176, + "D": 128, + "K": 4096, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.169311, + "evolution_flashlib_ms": 0.713215, + "evolution_tflops": 13.4764, + "evolution_flashlib_equiv_tflops": 3.1992, + "evolution_speedup": 4.2125, + }, + { + "source": "full95", + "label": "tail_odd_b1_n4224_k1024_d128", + "B": 1, + "N": 4224, + "D": 128, + "K": 1024, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.14784, + "evolution_flashlib_ms": 0.38112, + "evolution_tflops": 7.4898, + "evolution_flashlib_equiv_tflops": 2.9054, + "evolution_speedup": 2.5779, + }, + { + "source": "full95", + "label": "fallback_large_odd_b1_n8320_k1280_d128", + "B": 1, + "N": 8320, + "D": 128, + "K": 1280, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.149984, + "evolution_flashlib_ms": 0.345856, + "evolution_tflops": 18.1773, + "evolution_flashlib_equiv_tflops": 7.8828, + "evolution_speedup": 2.306, + }, + { + "source": "full95", + "label": "paired_large_b1_n65536_k4096_d128", + "B": 1, + "N": 65536, + "D": 128, + "K": 4096, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.259808, + "evolution_flashlib_ms": 1.096543, + "evolution_tflops": 264.501, + "evolution_flashlib_equiv_tflops": 62.6692, + "evolution_speedup": 4.2206, + }, + { + "source": "full95", + "label": "paired_large_n_b1_n131072_k256_d128", + "B": 1, + "N": 131072, + "D": 128, + "K": 256, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.184608, + "evolution_flashlib_ms": 0.271056, + "evolution_tflops": 46.5307, + "evolution_flashlib_equiv_tflops": 31.6907, + "evolution_speedup": 1.4683, + }, + { + "source": "full95", + "label": "paired_kover_b2_n256_k768_d128", + "B": 2, + "N": 256, + "D": 128, + "K": 768, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.16816, + "evolution_flashlib_ms": 0.630431, + "evolution_tflops": 0.5986, + "evolution_flashlib_equiv_tflops": 0.1597, + "evolution_speedup": 3.749, + }, + { + "source": "full95", + "label": "fallback_low_n_hugek_b2_n384_k8192_d128", + "B": 2, + "N": 384, + "D": 128, + "K": 8192, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.20048, + "evolution_flashlib_ms": 1.915198, + "evolution_tflops": 8.0338, + "evolution_flashlib_equiv_tflops": 0.841, + "evolution_speedup": 9.5531, + }, + { + "source": "full95", + "label": "paired_even_b2_n512_k8192_d128", + "B": 2, + "N": 512, + "D": 128, + "K": 8192, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.269792, + "evolution_flashlib_ms": 1.702303, + "evolution_tflops": 7.9598, + "evolution_flashlib_equiv_tflops": 1.2615, + "evolution_speedup": 6.3097, + }, + { + "source": "full95", + "label": "small_random_b2_n768_k512_d128", + "B": 2, + "N": 768, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "small_grid_single_tile_v10", + "evolution_kernel_ms": 0.144992, + "evolution_flashlib_ms": 0.426368, + "evolution_tflops": 1.3885, + "evolution_flashlib_equiv_tflops": 0.4722, + "evolution_speedup": 2.9406, + }, + { + "source": "full95", + "label": "random_b2_n768_k3072_d128", + "B": 2, + "N": 768, + "D": 128, + "K": 3072, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.319424, + "evolution_flashlib_ms": 0.817951, + "evolution_tflops": 3.7817, + "evolution_flashlib_equiv_tflops": 1.4768, + "evolution_speedup": 2.5607, + }, + { + "source": "full95", + "label": "smoke_b2_n1024_k512_d128", + "B": 2, + "N": 1024, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "small_grid_single_tile_v10", + "evolution_kernel_ms": 0.148896, + "evolution_flashlib_ms": 0.422031, + "evolution_tflops": 1.8028, + "evolution_flashlib_equiv_tflops": 0.6361, + "evolution_speedup": 2.8344, + }, + { + "source": "full95", + "label": "fallback_first_odd_b2_n1408_k256_d128", + "B": 2, + "N": 1408, + "D": 128, + "K": 256, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.14272, + "evolution_flashlib_ms": 0.280752, + "evolution_tflops": 1.2931, + "evolution_flashlib_equiv_tflops": 0.6573, + "evolution_speedup": 1.9672, + }, + { + "source": "full95", + "label": "paired_even_tail_b2_n1792_k256_d128", + "B": 2, + "N": 1792, + "D": 128, + "K": 256, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.164208, + "evolution_flashlib_ms": 0.248512, + "evolution_tflops": 1.4304, + "evolution_flashlib_equiv_tflops": 0.9451, + "evolution_speedup": 1.5134, + }, + { + "source": "full95", + "label": "tail_even_b2_n2048_k2048_d128", + "B": 2, + "N": 2048, + "D": 128, + "K": 2048, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.186303, + "evolution_flashlib_ms": 0.556255, + "evolution_tflops": 11.5268, + "evolution_flashlib_equiv_tflops": 3.8606, + "evolution_speedup": 2.9858, + }, + { + "source": "full95", + "label": "random_legal_b2_n4736_k512_d128", + "B": 2, + "N": 4736, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.146688, + "evolution_flashlib_ms": 0.346112, + "evolution_tflops": 8.4636, + "evolution_flashlib_equiv_tflops": 3.587, + "evolution_speedup": 2.3595, + }, + { + "source": "full95", + "label": "fallback_large_b2_n65664_k512_d128", + "B": 2, + "N": 65664, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.17792, + "evolution_flashlib_ms": 0.427904, + "evolution_tflops": 96.7478, + "evolution_flashlib_equiv_tflops": 40.2273, + "evolution_speedup": 2.405, + }, + { + "source": "full95", + "label": "small_low_b3_n256_k512_d128", + "B": 3, + "N": 256, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "small_grid_single_tile_v10", + "evolution_kernel_ms": 0.145791, + "evolution_flashlib_ms": 0.555391, + "evolution_tflops": 0.6905, + "evolution_flashlib_equiv_tflops": 0.1812, + "evolution_speedup": 3.8095, + }, + { + "source": "full95", + "label": "paired_first_even_b3_n1536_k512_d128", + "B": 3, + "N": 1536, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.166016, + "evolution_flashlib_ms": 0.352351, + "evolution_tflops": 3.6381, + "evolution_flashlib_equiv_tflops": 1.7141, + "evolution_speedup": 2.1224, + }, + { + "source": "full95", + "label": "fallback_mid_odd_b3_n2432_k512_d128", + "B": 3, + "N": 2432, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.144, + "evolution_flashlib_ms": 0.316095, + "evolution_tflops": 6.641, + "evolution_flashlib_equiv_tflops": 3.0254, + "evolution_speedup": 2.1951, + }, + { + "source": "full95", + "label": "random_legal_b3_n2560_k1536_d128", + "B": 3, + "N": 2560, + "D": 128, + "K": 1536, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.179008, + "evolution_flashlib_ms": 0.43128, + "evolution_tflops": 16.8702, + "evolution_flashlib_equiv_tflops": 7.0022, + "evolution_speedup": 2.4093, + }, + { + "source": "full95", + "label": "random_b3_n6144_k1280_d128", + "B": 3, + "N": 6144, + "D": 128, + "K": 1280, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.17632, + "evolution_flashlib_ms": 0.427359, + "evolution_tflops": 34.2548, + "evolution_flashlib_equiv_tflops": 14.1328, + "evolution_speedup": 2.4238, + }, + { + "source": "full95", + "label": "small_b4_n128_k512_d128", + "B": 4, + "N": 128, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "small_grid_single_tile_v10", + "evolution_kernel_ms": 0.144912, + "evolution_flashlib_ms": 0.484159, + "evolution_tflops": 0.4631, + "evolution_flashlib_equiv_tflops": 0.1386, + "evolution_speedup": 3.3411, + }, + { + "source": "full95", + "label": "random_legal_b4_n640_k1024_d128", + "B": 4, + "N": 640, + "D": 128, + "K": 1024, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.148095, + "evolution_flashlib_ms": 0.544863, + "evolution_tflops": 4.5315, + "evolution_flashlib_equiv_tflops": 1.2317, + "evolution_speedup": 3.6791, + }, + { + "source": "full95", + "label": "boundary_b4_n1024_k512_d128", + "B": 4, + "N": 1024, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "small_grid_single_tile_v10", + "evolution_kernel_ms": 0.144512, + "evolution_flashlib_ms": 0.394623, + "evolution_tflops": 3.7151, + "evolution_flashlib_equiv_tflops": 1.3605, + "evolution_speedup": 2.7307, + }, + { + "source": "full95", + "label": "fallback_odd_b4_n1664_k256_d128", + "B": 4, + "N": 1664, + "D": 128, + "K": 256, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.145728, + "evolution_flashlib_ms": 0.294912, + "evolution_tflops": 2.9933, + "evolution_flashlib_equiv_tflops": 1.4791, + "evolution_speedup": 2.0237, + }, + { + "source": "full95", + "label": "paired_widek_b4_n3072_k3072_d128", + "B": 4, + "N": 3072, + "D": 128, + "K": 3072, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.19712, + "evolution_flashlib_ms": 0.62648, + "evolution_tflops": 49.0242, + "evolution_flashlib_equiv_tflops": 15.4254, + "evolution_speedup": 3.1782, + }, + { + "source": "full95", + "label": "random_legal_b5_n512_k512_d128", + "B": 5, + "N": 512, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "small_grid_single_tile_v10", + "evolution_kernel_ms": 0.14368, + "evolution_flashlib_ms": 0.452576, + "evolution_tflops": 2.3354, + "evolution_flashlib_equiv_tflops": 0.7414, + "evolution_speedup": 3.1499, + }, + { + "source": "full95", + "label": "random_b5_n3456_k512_d128", + "B": 5, + "N": 3456, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.145872, + "evolution_flashlib_ms": 0.299296, + "evolution_tflops": 15.5268, + "evolution_flashlib_equiv_tflops": 7.5675, + "evolution_speedup": 2.0518, + }, + { + "source": "full95", + "label": "small_cap_b6_n1024_k256_d128", + "B": 6, + "N": 1024, + "D": 128, + "K": 256, + "dtype": "bfloat16", + "route": "small_grid_single_tile_v10", + "evolution_kernel_ms": 0.144144, + "evolution_flashlib_ms": 0.352896, + "evolution_tflops": 2.7934, + "evolution_flashlib_equiv_tflops": 1.141, + "evolution_speedup": 2.4482, + }, + { + "source": "full95", + "label": "paired_large_b6_n65536_k512_d128", + "B": 6, + "N": 65536, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.253025, + "evolution_flashlib_ms": 0.900256, + "evolution_tflops": 203.6937, + "evolution_flashlib_equiv_tflops": 57.2499, + "evolution_speedup": 3.558, + }, + { + "source": "full95", + "label": "small_boundary_b7_n896_k256_d128", + "B": 7, + "N": 896, + "D": 128, + "K": 256, + "dtype": "bfloat16", + "route": "small_grid_single_tile_v10", + "evolution_kernel_ms": 0.14368, + "evolution_flashlib_ms": 0.327696, + "evolution_tflops": 2.8608, + "evolution_flashlib_equiv_tflops": 1.2543, + "evolution_speedup": 2.2807, + }, + { + "source": "full95", + "label": "random_b7_n1024_k4096_d128", + "B": 7, + "N": 1024, + "D": 128, + "K": 4096, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.20928, + "evolution_flashlib_ms": 0.9272, + "evolution_tflops": 35.9145, + "evolution_flashlib_equiv_tflops": 8.1063, + "evolution_speedup": 4.4304, + }, + { + "source": "full95", + "label": "fallback_b8_n3712_k2048_d128", + "B": 8, + "N": 3712, + "D": 128, + "K": 2048, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.171536, + "evolution_flashlib_ms": 0.620543, + "evolution_tflops": 90.764, + "evolution_flashlib_equiv_tflops": 25.0897, + "evolution_speedup": 3.6176, + }, + { + "source": "full95", + "label": "paired_b8_n4096_k256_d128", + "B": 8, + "N": 4096, + "D": 128, + "K": 256, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.164624, + "evolution_flashlib_ms": 0.240416, + "evolution_tflops": 13.0448, + "evolution_flashlib_equiv_tflops": 8.9324, + "evolution_speedup": 1.4604, + }, + { + "source": "full95", + "label": "mid_b8_n8192_k512_d128", + "B": 8, + "N": 8192, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.176959, + "evolution_flashlib_ms": 0.395199, + "evolution_tflops": 48.5418, + "evolution_flashlib_equiv_tflops": 21.7357, + "evolution_speedup": 2.2333, + }, + { + "source": "full95", + "label": "random_legal_b9_n3840_k1792_d128", + "B": 9, + "N": 3840, + "D": 128, + "K": 1792, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.186464, + "evolution_flashlib_ms": 0.602848, + "evolution_tflops": 85.027, + "evolution_flashlib_equiv_tflops": 26.2993, + "evolution_speedup": 3.2331, + }, + { + "source": "full95", + "label": "small_b12_n640_k256_d128", + "B": 12, + "N": 640, + "D": 128, + "K": 256, + "dtype": "bfloat16", + "route": "small_grid_single_tile_v10", + "evolution_kernel_ms": 0.144128, + "evolution_flashlib_ms": 0.36936, + "evolution_tflops": 3.4921, + "evolution_flashlib_equiv_tflops": 1.3627, + "evolution_speedup": 2.5627, + }, + { + "source": "full95", + "label": "small_b16_n1024_k512_d128", + "B": 16, + "N": 1024, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "small_grid_single_tile_v10", + "evolution_kernel_ms": 0.1464, + "evolution_flashlib_ms": 0.40368, + "evolution_tflops": 14.6686, + "evolution_flashlib_equiv_tflops": 5.3198, + "evolution_speedup": 2.7574, + }, + { + "source": "full95", + "label": "large_b16_n32768_k1024_d128", + "B": 16, + "N": 32768, + "D": 128, + "K": 1024, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.363695, + "evolution_flashlib_ms": 2.084797, + "evolution_tflops": 377.8962, + "evolution_flashlib_equiv_tflops": 65.9244, + "evolution_speedup": 5.7323, + }, + { + "source": "full95", + "label": "paper_b32_n75776_k1024_d128", + "B": 32, + "N": 75776, + "D": 128, + "K": 1024, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.951326, + "evolution_flashlib_ms": 8.431702, + "evolution_tflops": 668.1781, + "evolution_flashlib_equiv_tflops": 75.3887, + "evolution_speedup": 8.8631, + }, + { + "source": "full95", + "label": "d160_paired_b1_n2048_k2048_d160", + "B": 1, + "N": 2048, + "D": 160, + "K": 2048, + "dtype": "bfloat16", + "route": "d160_padded_single_repeated_mma_v2", + "evolution_kernel_ms": 0.188448, + "evolution_flashlib_ms": 0.583535, + "evolution_tflops": 7.1223, + "evolution_flashlib_equiv_tflops": 2.3001, + "evolution_speedup": 3.0965, + }, + { + "source": "full95", + "label": "d160_fallback_b2_n2432_k1024_d160", + "B": 2, + "N": 2432, + "D": 160, + "K": 1024, + "dtype": "bfloat16", + "route": "d160_padded_single_repeated_mma_v2", + "evolution_kernel_ms": 0.18832, + "evolution_flashlib_ms": 0.419648, + "evolution_tflops": 8.4634, + "evolution_flashlib_equiv_tflops": 3.798, + "evolution_speedup": 2.2284, + }, + { + "source": "full95", + "label": "d192_fallback_b1_n2176_k1024_d192", + "B": 1, + "N": 2176, + "D": 192, + "K": 1024, + "dtype": "bfloat16", + "route": "d192_single_repeated_mma_v1", + "evolution_kernel_ms": 0.149952, + "evolution_flashlib_ms": 0.412832, + "evolution_tflops": 5.7061, + "evolution_flashlib_equiv_tflops": 2.0726, + "evolution_speedup": 2.7531, + }, + { + "source": "full95", + "label": "d192_small_b2_n1024_k512_d192", + "B": 2, + "N": 1024, + "D": 192, + "K": 512, + "dtype": "bfloat16", + "route": "d192_single_repeated_mma_v1", + "evolution_kernel_ms": 0.146657, + "evolution_flashlib_ms": 0.449999, + "evolution_tflops": 2.7455, + "evolution_flashlib_equiv_tflops": 0.8948, + "evolution_speedup": 3.0684, + }, + { + "source": "full95", + "label": "d192_paired_b2_n2048_k2048_d192", + "B": 2, + "N": 2048, + "D": 192, + "K": 2048, + "dtype": "bfloat16", + "route": "d192_paired_repeated_mma_v1", + "evolution_kernel_ms": 0.310336, + "evolution_flashlib_ms": 0.687759, + "evolution_tflops": 10.3798, + "evolution_flashlib_equiv_tflops": 4.6837, + "evolution_speedup": 2.2162, + }, + { + "source": "full95", + "label": "d192_large_b4_n32768_k1024_d192", + "B": 4, + "N": 32768, + "D": 192, + "K": 1024, + "dtype": "bfloat16", + "route": "d192_paired_repeated_mma_v1", + "evolution_kernel_ms": 0.221536, + "evolution_flashlib_ms": 0.913344, + "evolution_tflops": 232.6466, + "evolution_flashlib_equiv_tflops": 56.4296, + "evolution_speedup": 4.1228, + }, + { + "source": "full95", + "label": "d256_hugek_b1_n512_k8192_d256", + "B": 1, + "N": 512, + "D": 256, + "K": 8192, + "dtype": "bfloat16", + "route": "d256_single_repeated_mma_v1", + "evolution_kernel_ms": 0.222944, + "evolution_flashlib_ms": 55.197531, + "evolution_tflops": 9.6324, + "evolution_flashlib_equiv_tflops": 0.0389, + "evolution_speedup": 247.5847, + }, + { + "source": "full95", + "label": "d256_small_b1_n1024_k512_d256", + "B": 1, + "N": 1024, + "D": 256, + "K": 512, + "dtype": "bfloat16", + "route": "d256_single_repeated_mma_v1", + "evolution_kernel_ms": 0.148224, + "evolution_flashlib_ms": 0.429567, + "evolution_tflops": 1.811, + "evolution_flashlib_equiv_tflops": 0.6249, + "evolution_speedup": 2.8981, + }, + { + "source": "full95", + "label": "d256_paired_b1_n4096_k4096_d256", + "B": 1, + "N": 4096, + "D": 256, + "K": 4096, + "dtype": "bfloat16", + "route": "d256_single_repeated_mma_v1", + "evolution_kernel_ms": 0.179455, + "evolution_flashlib_ms": 55.152968, + "evolution_tflops": 47.8668, + "evolution_flashlib_equiv_tflops": 0.1557, + "evolution_speedup": 307.3359, + }, + { + "source": "full95", + "label": "d256_fallback_b2_n2432_k2048_d256", + "B": 2, + "N": 2432, + "D": 256, + "K": 2048, + "dtype": "bfloat16", + "route": "d256_single_repeated_mma_v1", + "evolution_kernel_ms": 0.16352, + "evolution_flashlib_ms": 0.616079, + "evolution_tflops": 31.1905, + "evolution_flashlib_equiv_tflops": 8.2786, + "evolution_speedup": 3.7676, + }, + { + "source": "full95", + "label": "d320_paired_b1_n2048_k4096_d320", + "B": 1, + "N": 2048, + "D": 320, + "K": 4096, + "dtype": "bfloat16", + "route": "highd_splitd_single_tile_6fcf_v1", + "evolution_kernel_ms": 0.23344, + "evolution_flashlib_ms": 46.53602, + "evolution_tflops": 22.9982, + "evolution_flashlib_equiv_tflops": 0.1154, + "evolution_speedup": 199.349, + }, + { + "source": "full95", + "label": "d320_large_b2_n16384_k1024_d320", + "B": 2, + "N": 16384, + "D": 320, + "K": 1024, + "dtype": "bfloat16", + "route": "highd_splitd_single_tile_6fcf_v1", + "evolution_kernel_ms": 0.19096, + "evolution_flashlib_ms": 0.586752, + "evolution_tflops": 112.4573, + "evolution_flashlib_equiv_tflops": 36.5995, + "evolution_speedup": 3.0726, + }, + { + "source": "full95", + "label": "d384_hugek_b1_n768_k8192_d384", + "B": 1, + "N": 768, + "D": 384, + "K": 8192, + "dtype": "bfloat16", + "route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "evolution_kernel_ms": 0.271664, + "evolution_flashlib_ms": 39.895303, + "evolution_tflops": 17.7861, + "evolution_flashlib_equiv_tflops": 0.1211, + "evolution_speedup": 146.8553, + }, + { + "source": "full95", + "label": "d384_small_b1_n896_k512_d384", + "B": 1, + "N": 896, + "D": 384, + "K": 512, + "dtype": "bfloat16", + "route": "highd_splitd_single_tile_6fcf_v1", + "evolution_kernel_ms": 0.15824, + "evolution_flashlib_ms": 0.521424, + "evolution_tflops": 2.2265, + "evolution_flashlib_equiv_tflops": 0.6757, + "evolution_speedup": 3.2951, + }, + { + "source": "full95", + "label": "d384_paired_b1_n2048_k4096_d384", + "B": 1, + "N": 2048, + "D": 384, + "K": 4096, + "dtype": "bfloat16", + "route": "highd_splitd_single_tile_6fcf_v1", + "evolution_kernel_ms": 0.252672, + "evolution_flashlib_ms": 38.477384, + "evolution_tflops": 25.4973, + "evolution_flashlib_equiv_tflops": 0.1674, + "evolution_speedup": 152.2819, + }, + { + "source": "full95", + "label": "d384_fallback_b3_n3456_k1024_d384", + "B": 3, + "N": 3456, + "D": 384, + "K": 1024, + "dtype": "bfloat16", + "route": "highd_splitd_single_tile_6fcf_v1", + "evolution_kernel_ms": 0.174575, + "evolution_flashlib_ms": 0.595328, + "evolution_tflops": 46.706, + "evolution_flashlib_equiv_tflops": 13.6962, + "evolution_speedup": 3.4101, + }, + { + "source": "full95", + "label": "d448_hugek_b1_n512_k8192_d448", + "B": 1, + "N": 512, + "D": 448, + "K": 8192, + "dtype": "bfloat16", + "route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "evolution_kernel_ms": 0.262959, + "evolution_flashlib_ms": 40.698431, + "evolution_tflops": 14.2916, + "evolution_flashlib_equiv_tflops": 0.0923, + "evolution_speedup": 154.771, + }, + { + "source": "full95", + "label": "d448_paired_b1_n2048_k4096_d448", + "B": 1, + "N": 2048, + "D": 448, + "K": 4096, + "dtype": "bfloat16", + "route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "evolution_kernel_ms": 0.269632, + "evolution_flashlib_ms": 38.67485, + "evolution_tflops": 27.8757, + "evolution_flashlib_equiv_tflops": 0.1943, + "evolution_speedup": 143.4357, + }, + { + "source": "full95", + "label": "d512_hugek_b1_n512_k8192_d512", + "B": 1, + "N": 512, + "D": 512, + "K": 8192, + "dtype": "bfloat16", + "route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "evolution_kernel_ms": 0.266288, + "evolution_flashlib_ms": 32.378557, + "evolution_tflops": 16.1291, + "evolution_flashlib_equiv_tflops": 0.1326, + "evolution_speedup": 121.5925, + }, + { + "source": "full95", + "label": "d512_small_b1_n1024_k512_d512", + "B": 1, + "N": 1024, + "D": 512, + "K": 512, + "dtype": "bfloat16", + "route": "highd_splitd_single_tile_6fcf_v1", + "evolution_kernel_ms": 0.162688, + "evolution_flashlib_ms": 0.554048, + "evolution_tflops": 3.3, + "evolution_flashlib_equiv_tflops": 0.969, + "evolution_speedup": 3.4056, + }, + { + "source": "full95", + "label": "d512_paired_b1_n2048_k4096_d512", + "B": 1, + "N": 2048, + "D": 512, + "K": 4096, + "dtype": "bfloat16", + "route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "evolution_kernel_ms": 0.270976, + "evolution_flashlib_ms": 30.573365, + "evolution_tflops": 31.7, + "evolution_flashlib_equiv_tflops": 0.281, + "evolution_speedup": 112.8268, + }, + { + "source": "full95", + "label": "d512_fallback_b2_n2432_k2048_d512", + "B": 2, + "N": 2432, + "D": 512, + "K": 2048, + "dtype": "bfloat16", + "route": "highd_splitd_single_tile_6fcf_v1", + "evolution_kernel_ms": 0.220319, + "evolution_flashlib_ms": 0.855488, + "evolution_tflops": 46.299, + "evolution_flashlib_equiv_tflops": 11.9237, + "evolution_speedup": 3.883, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d16_b4_n32768_k1024_d16", + "B": 4, + "N": 32768, + "D": 16, + "K": 1024, + "dtype": "bfloat16", + "route": "microdim_hybrid_9c0d_v1", + "evolution_kernel_ms": 0.216591, + "evolution_flashlib_ms": 0.330559, + "evolution_tflops": 19.8298, + "evolution_flashlib_equiv_tflops": 12.993, + "evolution_speedup": 1.5262, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d16_b8_n65536_k512_d16", + "B": 8, + "N": 65536, + "D": 16, + "K": 512, + "dtype": "bfloat16", + "route": "microdim_hybrid_9c0d_v1", + "evolution_kernel_ms": 0.295359, + "evolution_flashlib_ms": 0.444799, + "evolution_tflops": 29.083, + "evolution_flashlib_equiv_tflops": 19.3119, + "evolution_speedup": 1.506, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d32_b4_n32768_k1024_d32", + "B": 4, + "N": 32768, + "D": 32, + "K": 1024, + "dtype": "bfloat16", + "route": "microdim_hybrid_9c0d_v1", + "evolution_kernel_ms": 0.218816, + "evolution_flashlib_ms": 0.360896, + "evolution_tflops": 39.2564, + "evolution_flashlib_equiv_tflops": 23.8017, + "evolution_speedup": 1.6493, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d32_b8_n65536_k512_d32", + "B": 8, + "N": 65536, + "D": 32, + "K": 512, + "dtype": "bfloat16", + "route": "microdim_hybrid_9c0d_v1", + "evolution_kernel_ms": 0.301439, + "evolution_flashlib_ms": 0.500015, + "evolution_tflops": 56.9928, + "evolution_flashlib_equiv_tflops": 34.3587, + "evolution_speedup": 1.6588, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_3328_d48_small_boundary_b1_n256_k256_d48", + "B": 1, + "N": 256, + "D": 48, + "K": 256, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.168656, + "evolution_flashlib_ms": 0.434063, + "evolution_tflops": 0.0373, + "evolution_flashlib_equiv_tflops": 0.0145, + "evolution_speedup": 2.5737, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d48_b1_n512_k8192_d48", + "B": 1, + "N": 512, + "D": 48, + "K": 8192, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.21952, + "evolution_flashlib_ms": 1.464301, + "evolution_tflops": 1.8342, + "evolution_flashlib_equiv_tflops": 0.275, + "evolution_speedup": 6.6705, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_d9d5_d48_highk_heldout_b2_n640_k4096_d48", + "B": 2, + "N": 640, + "D": 48, + "K": 4096, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.19392, + "evolution_flashlib_ms": 0.916111, + "evolution_tflops": 2.5955, + "evolution_flashlib_equiv_tflops": 0.5494, + "evolution_speedup": 4.7242, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_a2f8_d48_k1024_boundary_b2_n768_k1024_d48", + "B": 2, + "N": 768, + "D": 48, + "K": 1024, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.173184, + "evolution_flashlib_ms": 0.497055, + "evolution_tflops": 0.8719, + "evolution_flashlib_equiv_tflops": 0.3038, + "evolution_speedup": 2.8701, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_8f09_d48_medium_tail_b2_n1792_k512_d48", + "B": 2, + "N": 1792, + "D": 48, + "K": 512, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.169375, + "evolution_flashlib_ms": 0.316848, + "evolution_tflops": 1.0401, + "evolution_flashlib_equiv_tflops": 0.556, + "evolution_speedup": 1.8707, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d48_b2_n2048_k512_d48", + "B": 2, + "N": 2048, + "D": 48, + "K": 512, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.166656, + "evolution_flashlib_ms": 0.30624, + "evolution_tflops": 1.208, + "evolution_flashlib_equiv_tflops": 0.6574, + "evolution_speedup": 1.8376, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_5600_d48_low_n_boundary_b3_n1536_k256_d48", + "B": 3, + "N": 1536, + "D": 48, + "K": 256, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.169183, + "evolution_flashlib_ms": 0.281087, + "evolution_tflops": 0.6694, + "evolution_flashlib_equiv_tflops": 0.4029, + "evolution_speedup": 1.6614, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_c44f_d48_midk_boundary_b3_n2688_k768_d48", + "B": 3, + "N": 2688, + "D": 48, + "K": 768, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.174383, + "evolution_flashlib_ms": 0.309952, + "evolution_tflops": 3.4094, + "evolution_flashlib_equiv_tflops": 1.9182, + "evolution_speedup": 1.7774, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_ef00_d48_midn_boundary_b3_n3456_k512_d48", + "B": 3, + "N": 3456, + "D": 48, + "K": 512, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.172495, + "evolution_flashlib_ms": 0.257807, + "evolution_tflops": 2.9543, + "evolution_flashlib_equiv_tflops": 1.9767, + "evolution_speedup": 1.4946, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_68cf_d48_boundary_b4_n2304_k512_d48", + "B": 4, + "N": 2304, + "D": 48, + "K": 512, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.170608, + "evolution_flashlib_ms": 0.293984, + "evolution_tflops": 2.6551, + "evolution_flashlib_equiv_tflops": 1.5408, + "evolution_speedup": 1.7232, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_1d49_d48_lowk_boundary_b4_n3968_k256_d48", + "B": 4, + "N": 3968, + "D": 48, + "K": 256, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.17192, + "evolution_flashlib_ms": 0.228768, + "evolution_tflops": 2.2689, + "evolution_flashlib_equiv_tflops": 1.7051, + "evolution_speedup": 1.3307, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d48_b4_n8192_k1024_d48", + "B": 4, + "N": 8192, + "D": 48, + "K": 1024, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.179807, + "evolution_flashlib_ms": 0.335967, + "evolution_tflops": 17.9149, + "evolution_flashlib_equiv_tflops": 9.5879, + "evolution_speedup": 1.8685, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_9ca0_d48_guard_boundary_b6_n12288_k512_d48", + "B": 6, + "N": 12288, + "D": 48, + "K": 512, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.182336, + "evolution_flashlib_ms": 0.318304, + "evolution_tflops": 19.8747, + "evolution_flashlib_equiv_tflops": 11.385, + "evolution_speedup": 1.7457, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d112_b1_n256_k256_d112", + "B": 1, + "N": 256, + "D": 112, + "K": 256, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.166575, + "evolution_flashlib_ms": 0.416031, + "evolution_tflops": 0.0881, + "evolution_flashlib_equiv_tflops": 0.0353, + "evolution_speedup": 2.4976, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_8f09_d112_small_highk_b1_n384_k4096_d112", + "B": 1, + "N": 384, + "D": 112, + "K": 4096, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.195296, + "evolution_flashlib_ms": 1.266206, + "evolution_tflops": 1.804, + "evolution_flashlib_equiv_tflops": 0.2782, + "evolution_speedup": 6.4835, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d112_b1_n512_k8192_d112", + "B": 1, + "N": 512, + "D": 112, + "K": 8192, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.294304, + "evolution_flashlib_ms": 1.731693, + "evolution_tflops": 3.1924, + "evolution_flashlib_equiv_tflops": 0.5425, + "evolution_speedup": 5.884, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_1d49_d112_highk_tail_b1_n1408_k4096_d112", + "B": 1, + "N": 1408, + "D": 112, + "K": 4096, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.196192, + "evolution_flashlib_ms": 0.866559, + "evolution_tflops": 6.5846, + "evolution_flashlib_equiv_tflops": 1.4908, + "evolution_speedup": 4.4169, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_3328_d112_highk_low_n_b2_n512_k8192_d112", + "B": 2, + "N": 512, + "D": 112, + "K": 8192, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.295872, + "evolution_flashlib_ms": 1.758429, + "evolution_tflops": 6.3509, + "evolution_flashlib_equiv_tflops": 1.0686, + "evolution_speedup": 5.9432, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d112_b2_n2048_k512_d112", + "B": 2, + "N": 2048, + "D": 112, + "K": 512, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.188, + "evolution_flashlib_ms": 0.312623, + "evolution_tflops": 2.4987, + "evolution_flashlib_equiv_tflops": 1.5026, + "evolution_speedup": 1.6629, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_c44f_d112_odd_tail_b2_n3200_k768_d112", + "B": 2, + "N": 3200, + "D": 112, + "K": 768, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.173088, + "evolution_flashlib_ms": 0.321696, + "evolution_tflops": 6.361, + "evolution_flashlib_equiv_tflops": 3.4225, + "evolution_speedup": 1.8586, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_d9d5_d112_highk_request_b3_n768_k8192_d112", + "B": 3, + "N": 768, + "D": 112, + "K": 8192, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.287872, + "evolution_flashlib_ms": 1.518318, + "evolution_tflops": 14.6866, + "evolution_flashlib_equiv_tflops": 2.7846, + "evolution_speedup": 5.2743, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_9ca0_d112_tail_div_b3_n3840_k768_d112", + "B": 3, + "N": 3840, + "D": 112, + "K": 768, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.194592, + "evolution_flashlib_ms": 0.333232, + "evolution_tflops": 10.1844, + "evolution_flashlib_equiv_tflops": 5.9472, + "evolution_speedup": 1.7125, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_a2f8_d112_lowk_tail_b4_n3456_k256_d112", + "B": 4, + "N": 3456, + "D": 112, + "K": 256, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.168912, + "evolution_flashlib_ms": 0.231359, + "evolution_tflops": 4.6931, + "evolution_flashlib_equiv_tflops": 3.4264, + "evolution_speedup": 1.3697, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_ef00_d112_odd_tail_b4_n3712_k1024_d112", + "B": 4, + "N": 3712, + "D": 112, + "K": 1024, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.176608, + "evolution_flashlib_ms": 0.386799, + "evolution_tflops": 19.2844, + "evolution_flashlib_equiv_tflops": 8.805, + "evolution_speedup": 2.1902, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d112_b4_n8192_k1024_d112", + "B": 4, + "N": 8192, + "D": 112, + "K": 1024, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.1956, + "evolution_flashlib_ms": 0.415103, + "evolution_tflops": 38.4264, + "evolution_flashlib_equiv_tflops": 18.1068, + "evolution_speedup": 2.1222, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_5600_d112_odd_tile_tail_b5_n2176_k512_d112", + "B": 5, + "N": 2176, + "D": 112, + "K": 512, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.171519, + "evolution_flashlib_ms": 0.338463, + "evolution_tflops": 7.275, + "evolution_flashlib_equiv_tflops": 3.6867, + "evolution_speedup": 1.9733, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_68cf_d112_tail_b5_n2944_k512_d112", + "B": 5, + "N": 2944, + "D": 112, + "K": 512, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.173023, + "evolution_flashlib_ms": 0.296095, + "evolution_tflops": 9.7571, + "evolution_flashlib_equiv_tflops": 5.7016, + "evolution_speedup": 1.7113, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d128_paired_b2_n262144_k256_d128", + "B": 2, + "N": 262144, + "D": 128, + "K": 256, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.24464, + "evolution_flashlib_ms": 0.663294, + "evolution_tflops": 140.4505, + "evolution_flashlib_equiv_tflops": 51.8017, + "evolution_speedup": 2.7113, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d128_fallback_b3_n1920_k256_d128", + "B": 3, + "N": 1920, + "D": 128, + "K": 256, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.1436, + "evolution_flashlib_ms": 0.281664, + "evolution_tflops": 2.6287, + "evolution_flashlib_equiv_tflops": 1.3402, + "evolution_speedup": 1.9614, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_8f09_d128_fallback_neighbor_b4_n4480_k512_d128", + "B": 4, + "N": 4480, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.145856, + "evolution_flashlib_ms": 0.325535, + "evolution_tflops": 16.1036, + "evolution_flashlib_equiv_tflops": 7.2152, + "evolution_speedup": 2.2319, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_ef00_d128_forced_fallback_b4_n7552_k512_d128", + "B": 4, + "N": 7552, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.153023, + "evolution_flashlib_ms": 0.314032, + "evolution_tflops": 25.8747, + "evolution_flashlib_equiv_tflops": 12.6084, + "evolution_speedup": 2.0522, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d128_fallback_b5_n2176_k512_d128", + "B": 5, + "N": 2176, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.14432, + "evolution_flashlib_ms": 0.309232, + "evolution_tflops": 9.8813, + "evolution_flashlib_equiv_tflops": 4.6116, + "evolution_speedup": 2.1427, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_d9d5_d128_guard_miss_fallback_b5_n6016_k1024_d128", + "B": 5, + "N": 6016, + "D": 128, + "K": 1024, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.159104, + "evolution_flashlib_ms": 0.460207, + "evolution_tflops": 49.5606, + "evolution_flashlib_equiv_tflops": 17.1342, + "evolution_speedup": 2.8925, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_1d49_d128_forced_fallback_b5_n7296_k512_d128", + "B": 5, + "N": 7296, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.152639, + "evolution_flashlib_ms": 0.339855, + "evolution_tflops": 31.3255, + "evolution_flashlib_equiv_tflops": 14.0693, + "evolution_speedup": 2.2265, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_c44f_d128_forced_fallback_b6_n6272_k512_d128", + "B": 6, + "N": 6272, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.152256, + "evolution_flashlib_ms": 0.333376, + "evolution_tflops": 32.3961, + "evolution_flashlib_equiv_tflops": 14.7956, + "evolution_speedup": 2.1896, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_a2f8_d128_odd_fallback_b6_n8576_k512_d128", + "B": 6, + "N": 8576, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.158688, + "evolution_flashlib_ms": 0.349887, + "evolution_tflops": 42.5013, + "evolution_flashlib_equiv_tflops": 19.2761, + "evolution_speedup": 2.2049, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d128_fallback_b7_n2432_k1024_d128", + "B": 7, + "N": 2432, + "D": 128, + "K": 1024, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.147936, + "evolution_flashlib_ms": 0.4056, + "evolution_tflops": 30.1667, + "evolution_flashlib_equiv_tflops": 11.0028, + "evolution_speedup": 2.7417, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_68cf_d128_forced_fallback_b7_n6016_k512_d128", + "B": 7, + "N": 6016, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.156832, + "evolution_flashlib_ms": 0.386527, + "evolution_tflops": 35.195, + "evolution_flashlib_equiv_tflops": 14.2803, + "evolution_speedup": 2.4646, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_9ca0_d128_exact_guard_neighbor_b8_n8064_k256_d128", + "B": 8, + "N": 8064, + "D": 128, + "K": 256, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.155168, + "evolution_flashlib_ms": 0.291264, + "evolution_tflops": 27.247, + "evolution_flashlib_equiv_tflops": 14.5156, + "evolution_speedup": 1.8771, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d128_paired_b8_n8192_k256_d128", + "B": 8, + "N": 8192, + "D": 128, + "K": 256, + "dtype": "bfloat16", + "route": "d128_even_near_floor_v10_repair", + "evolution_kernel_ms": 0.15424, + "evolution_flashlib_ms": 0.292671, + "evolution_tflops": 27.846, + "evolution_flashlib_equiv_tflops": 14.6751, + "evolution_speedup": 1.8975, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_3328_d128_exact_guard_k_neighbor_b8_n8192_k512_d128", + "B": 8, + "N": 8192, + "D": 128, + "K": 512, + "dtype": "bfloat16", + "route": "paired_large_v15", + "evolution_kernel_ms": 0.177472, + "evolution_flashlib_ms": 0.394592, + "evolution_tflops": 48.4016, + "evolution_flashlib_equiv_tflops": 21.7692, + "evolution_speedup": 2.2234, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_5600_d128_fallback_neighbor_b8_n8320_k256_d128", + "B": 8, + "N": 8320, + "D": 128, + "K": 256, + "dtype": "bfloat16", + "route": "aligned_v10_fallback", + "evolution_kernel_ms": 0.155872, + "evolution_flashlib_ms": 0.297119, + "evolution_tflops": 27.985, + "evolution_flashlib_equiv_tflops": 14.6812, + "evolution_speedup": 1.9062, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d144_b1_n256_k256_d144", + "B": 1, + "N": 256, + "D": 144, + "K": 256, + "dtype": "bfloat16", + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_kernel_ms": 0.167728, + "evolution_flashlib_ms": 0.537279, + "evolution_tflops": 0.1125, + "evolution_flashlib_equiv_tflops": 0.0351, + "evolution_speedup": 3.2033, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d144_b1_n512_k8192_d144", + "B": 1, + "N": 512, + "D": 144, + "K": 8192, + "dtype": "bfloat16", + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_kernel_ms": 0.302528, + "evolution_flashlib_ms": 2.001693, + "evolution_tflops": 3.9929, + "evolution_flashlib_equiv_tflops": 0.6035, + "evolution_speedup": 6.6166, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d144_b2_n2048_k1024_d144", + "B": 2, + "N": 2048, + "D": 144, + "K": 1024, + "dtype": "bfloat16", + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_kernel_ms": 0.192896, + "evolution_flashlib_ms": 0.465663, + "evolution_tflops": 6.2622, + "evolution_flashlib_equiv_tflops": 2.5941, + "evolution_speedup": 2.4141, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d144_b4_n8192_k1024_d144", + "B": 4, + "N": 8192, + "D": 144, + "K": 1024, + "dtype": "bfloat16", + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_kernel_ms": 0.195104, + "evolution_flashlib_ms": 0.473791, + "evolution_tflops": 49.5309, + "evolution_flashlib_equiv_tflops": 20.3965, + "evolution_speedup": 2.4284, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d176_b1_n256_k256_d176", + "B": 1, + "N": 256, + "D": 176, + "K": 256, + "dtype": "bfloat16", + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_kernel_ms": 0.169792, + "evolution_flashlib_ms": 0.509327, + "evolution_tflops": 0.1359, + "evolution_flashlib_equiv_tflops": 0.0453, + "evolution_speedup": 2.9997, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d176_b1_n512_k8192_d176", + "B": 1, + "N": 512, + "D": 176, + "K": 8192, + "dtype": "bfloat16", + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_kernel_ms": 0.301759, + "evolution_flashlib_ms": 2.050605, + "evolution_tflops": 4.8926, + "evolution_flashlib_equiv_tflops": 0.72, + "evolution_speedup": 6.7955, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d176_b2_n2048_k1024_d176", + "B": 2, + "N": 2048, + "D": 176, + "K": 1024, + "dtype": "bfloat16", + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_kernel_ms": 0.191839, + "evolution_flashlib_ms": 0.486367, + "evolution_tflops": 7.696, + "evolution_flashlib_equiv_tflops": 3.0356, + "evolution_speedup": 2.5353, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d176_b4_n8192_k1024_d176", + "B": 4, + "N": 8192, + "D": 176, + "K": 1024, + "dtype": "bfloat16", + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_kernel_ms": 0.195712, + "evolution_flashlib_ms": 0.492175, + "evolution_tflops": 60.3497, + "evolution_flashlib_equiv_tflops": 23.9979, + "evolution_speedup": 2.5148, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d224_b1_n256_k256_d224", + "B": 1, + "N": 256, + "D": 224, + "K": 256, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.170496, + "evolution_flashlib_ms": 0.480639, + "evolution_tflops": 0.1722, + "evolution_flashlib_equiv_tflops": 0.0611, + "evolution_speedup": 2.8191, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d224_b1_n512_k8192_d224", + "B": 1, + "N": 512, + "D": 224, + "K": 8192, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.236528, + "evolution_flashlib_ms": 2.069245, + "evolution_tflops": 7.9443, + "evolution_flashlib_equiv_tflops": 0.9081, + "evolution_speedup": 8.7484, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d224_b2_n2048_k1024_d224", + "B": 2, + "N": 2048, + "D": 224, + "K": 1024, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.176704, + "evolution_flashlib_ms": 0.47112, + "evolution_tflops": 10.6339, + "evolution_flashlib_equiv_tflops": 3.9885, + "evolution_speedup": 2.6662, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_d9d5_d224_random_midk_b2_n2944_k1280_d224", + "B": 2, + "N": 2944, + "D": 224, + "K": 1280, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.180384, + "evolution_flashlib_ms": 0.488255, + "evolution_tflops": 18.7179, + "evolution_flashlib_equiv_tflops": 6.9153, + "evolution_speedup": 2.7068, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_9ca0_d224_guard_overlap_b2_n4096_k256_d224", + "B": 2, + "N": 4096, + "D": 224, + "K": 256, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.1712, + "evolution_flashlib_ms": 0.258143, + "evolution_tflops": 5.4879, + "evolution_flashlib_equiv_tflops": 3.6395, + "evolution_speedup": 1.5078, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_ef00_d224_midn_overlap_b2_n5376_k512_d224", + "B": 2, + "N": 5376, + "D": 224, + "K": 512, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.177808, + "evolution_flashlib_ms": 0.36856, + "evolution_tflops": 13.8703, + "evolution_flashlib_equiv_tflops": 6.6916, + "evolution_speedup": 2.0728, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_a2f8_d224_highn_overlap_b2_n6144_k768_d224", + "B": 2, + "N": 6144, + "D": 224, + "K": 768, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.176256, + "evolution_flashlib_ms": 0.45224, + "evolution_tflops": 23.987, + "evolution_flashlib_equiv_tflops": 9.3487, + "evolution_speedup": 2.5658, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_5600_d224_k512_overlap_b3_n3072_k512_d224", + "B": 3, + "N": 3072, + "D": 224, + "K": 512, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.174928, + "evolution_flashlib_ms": 0.343839, + "evolution_tflops": 12.0846, + "evolution_flashlib_equiv_tflops": 6.148, + "evolution_speedup": 1.9656, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_3328_d224_tail_div_b3_n3840_k512_d224", + "B": 3, + "N": 3840, + "D": 224, + "K": 512, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.174895, + "evolution_flashlib_ms": 0.318207, + "evolution_tflops": 15.1086, + "evolution_flashlib_equiv_tflops": 8.3041, + "evolution_speedup": 1.8194, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_68cf_d224_overlap_b3_n5120_k1024_d224", + "B": 3, + "N": 5120, + "D": 224, + "K": 1024, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.180032, + "evolution_flashlib_ms": 0.486208, + "evolution_tflops": 39.1399, + "evolution_flashlib_equiv_tflops": 14.4926, + "evolution_speedup": 2.7007, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_8f09_d224_low_n_boundary_b4_n1536_k256_d224", + "B": 4, + "N": 1536, + "D": 224, + "K": 256, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.173664, + "evolution_flashlib_ms": 0.317296, + "evolution_tflops": 4.0575, + "evolution_flashlib_equiv_tflops": 2.2208, + "evolution_speedup": 1.8271, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_c44f_d224_overlap_b4_n4480_k512_d224", + "B": 4, + "N": 4480, + "D": 224, + "K": 512, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.173008, + "evolution_flashlib_ms": 0.402944, + "evolution_tflops": 23.7585, + "evolution_flashlib_equiv_tflops": 10.201, + "evolution_speedup": 2.329, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d224_b4_n8192_k1024_d224", + "B": 4, + "N": 8192, + "D": 224, + "K": 1024, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.190783, + "evolution_flashlib_ms": 0.553311, + "evolution_tflops": 78.7929, + "evolution_flashlib_equiv_tflops": 27.1681, + "evolution_speedup": 2.9002, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_1d49_d224_highn_overlap_b5_n5632_k768_d224", + "B": 5, + "N": 5632, + "D": 224, + "K": 768, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.186336, + "evolution_flashlib_ms": 0.465423, + "evolution_tflops": 51.9966, + "evolution_flashlib_equiv_tflops": 20.8173, + "evolution_speedup": 2.4978, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d288_b1_n256_k256_d288", + "B": 1, + "N": 256, + "D": 288, + "K": 256, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.171104, + "evolution_flashlib_ms": 0.581567, + "evolution_tflops": 0.2206, + "evolution_flashlib_equiv_tflops": 0.0649, + "evolution_speedup": 3.3989, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_9ca0_d288_highk_low_n_b1_n384_k4096_d288", + "B": 1, + "N": 384, + "D": 288, + "K": 4096, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.243872, + "evolution_flashlib_ms": 56.807529, + "evolution_tflops": 3.7149, + "evolution_flashlib_equiv_tflops": 0.0159, + "evolution_speedup": 232.9399, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d288_b1_n512_k8192_d288", + "B": 1, + "N": 512, + "D": 288, + "K": 8192, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.273615, + "evolution_flashlib_ms": 57.444236, + "evolution_tflops": 8.8296, + "evolution_flashlib_equiv_tflops": 0.0421, + "evolution_speedup": 209.9455, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_a2f8_d288_splitk_probe_b1_n640_k4096_d288", + "B": 1, + "N": 640, + "D": 288, + "K": 4096, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.24416, + "evolution_flashlib_ms": 57.6586, + "evolution_tflops": 6.1843, + "evolution_flashlib_equiv_tflops": 0.0262, + "evolution_speedup": 236.1509, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_d9d5_d288_heldout_low_n_b1_n896_k4096_d288", + "B": 1, + "N": 896, + "D": 288, + "K": 4096, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.2448, + "evolution_flashlib_ms": 57.939113, + "evolution_tflops": 8.6353, + "evolution_flashlib_equiv_tflops": 0.0365, + "evolution_speedup": 236.6794, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_ef00_d288_highk_heldout_b1_n1664_k4096_d288", + "B": 1, + "N": 1664, + "D": 288, + "K": 4096, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.247408, + "evolution_flashlib_ms": 58.88955, + "evolution_tflops": 15.868, + "evolution_flashlib_equiv_tflops": 0.0667, + "evolution_speedup": 238.0261, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_8f09_d288_k768_overlap_b1_n2560_k768_d288", + "B": 1, + "N": 2560, + "D": 288, + "K": 768, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.18432, + "evolution_flashlib_ms": 0.480287, + "evolution_tflops": 6.144, + "evolution_flashlib_equiv_tflops": 2.3579, + "evolution_speedup": 2.6057, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_c44f_d288_highk_heldout_b2_n768_k4096_d288", + "B": 2, + "N": 768, + "D": 288, + "K": 4096, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.246496, + "evolution_flashlib_ms": 1.583007, + "evolution_tflops": 14.7016, + "evolution_flashlib_equiv_tflops": 2.2892, + "evolution_speedup": 6.422, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_5600_d288_mid_highk_b2_n1024_k2048_d288", + "B": 2, + "N": 1024, + "D": 288, + "K": 2048, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.206512, + "evolution_flashlib_ms": 0.976096, + "evolution_tflops": 11.6987, + "evolution_flashlib_equiv_tflops": 2.4751, + "evolution_speedup": 4.7266, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_68cf_d288_boundary_b2_n1920_k512_d288", + "B": 2, + "N": 1920, + "D": 288, + "K": 512, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.177824, + "evolution_flashlib_ms": 0.421007, + "evolution_tflops": 6.3684, + "evolution_flashlib_equiv_tflops": 2.6899, + "evolution_speedup": 2.3676, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d288_b2_n2048_k1024_d288", + "B": 2, + "N": 2048, + "D": 288, + "K": 1024, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.18616, + "evolution_flashlib_ms": 0.65016, + "evolution_tflops": 12.9777, + "evolution_flashlib_equiv_tflops": 3.7159, + "evolution_speedup": 3.4925, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_1d49_d288_low_n_heldout_b3_n1152_k2048_d288", + "B": 3, + "N": 1152, + "D": 288, + "K": 2048, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.207744, + "evolution_flashlib_ms": 1.007006, + "evolution_tflops": 19.6245, + "evolution_flashlib_equiv_tflops": 4.0485, + "evolution_speedup": 4.8473, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_3328_d288_guard_overlap_b4_n8192_k256_d288", + "B": 4, + "N": 8192, + "D": 288, + "K": 256, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.182432, + "evolution_flashlib_ms": 0.344959, + "evolution_tflops": 26.4857, + "evolution_flashlib_equiv_tflops": 14.007, + "evolution_speedup": 1.8909, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d288_b4_n8192_k1024_d288", + "B": 4, + "N": 8192, + "D": 288, + "K": 1024, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.214687, + "evolution_flashlib_ms": 0.669808, + "evolution_tflops": 90.0257, + "evolution_flashlib_equiv_tflops": 28.8551, + "evolution_speedup": 3.1199, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d352_b1_n256_k256_d352", + "B": 1, + "N": 256, + "D": 352, + "K": 256, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.172703, + "evolution_flashlib_ms": 0.627167, + "evolution_tflops": 0.2671, + "evolution_flashlib_equiv_tflops": 0.0736, + "evolution_speedup": 3.6315, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_9ca0_d352_splitk_boundary_b1_n512_k8192_d352", + "B": 1, + "N": 512, + "D": 352, + "K": 8192, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.275519, + "evolution_flashlib_ms": 49.145822, + "evolution_tflops": 10.7172, + "evolution_flashlib_equiv_tflops": 0.0601, + "evolution_speedup": 178.3754, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d352_b1_n512_k8192_d352", + "B": 1, + "N": 512, + "D": 352, + "K": 8192, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.272767, + "evolution_flashlib_ms": 47.910581, + "evolution_tflops": 10.8253, + "evolution_flashlib_equiv_tflops": 0.0616, + "evolution_speedup": 175.6465, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_c44f_d352_random_b1_n3328_k768_d352", + "B": 1, + "N": 3328, + "D": 352, + "K": 768, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.187632, + "evolution_flashlib_ms": 0.4864, + "evolution_tflops": 9.5898, + "evolution_flashlib_equiv_tflops": 3.6993, + "evolution_speedup": 2.5923, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_5600_d352_splitk_edge_b2_n768_k4096_d352", + "B": 2, + "N": 768, + "D": 352, + "K": 4096, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 1.051615, + "evolution_flashlib_ms": 1.621022, + "evolution_tflops": 4.2118, + "evolution_flashlib_equiv_tflops": 2.7323, + "evolution_speedup": 1.5415, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_1d49_d352_request_highk_b2_n1024_k8192_d352", + "B": 2, + "N": 1024, + "D": 352, + "K": 8192, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.302655, + "evolution_flashlib_ms": 2.602461, + "evolution_tflops": 39.0252, + "evolution_flashlib_equiv_tflops": 4.5385, + "evolution_speedup": 8.5988, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d352_b2_n2048_k1024_d352", + "B": 2, + "N": 2048, + "D": 352, + "K": 1024, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.189184, + "evolution_flashlib_ms": 0.604831, + "evolution_tflops": 15.608, + "evolution_flashlib_equiv_tflops": 4.882, + "evolution_speedup": 3.1971, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_ef00_d352_request_low_n_b3_n896_k8192_d352", + "B": 3, + "N": 896, + "D": 352, + "K": 8192, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.316191, + "evolution_flashlib_ms": 3.136652, + "evolution_tflops": 49.0278, + "evolution_flashlib_equiv_tflops": 4.9423, + "evolution_speedup": 9.9201, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_3328_d352_random_legal_b3_n2048_k768_d352", + "B": 3, + "N": 2048, + "D": 352, + "K": 768, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.183872, + "evolution_flashlib_ms": 0.582335, + "evolution_tflops": 18.0663, + "evolution_flashlib_equiv_tflops": 5.7044, + "evolution_speedup": 3.1671, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_68cf_d352_tail_b3_n2816_k768_d352", + "B": 3, + "N": 2816, + "D": 352, + "K": 768, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.18768, + "evolution_flashlib_ms": 0.500207, + "evolution_tflops": 24.3372, + "evolution_flashlib_equiv_tflops": 9.1314, + "evolution_speedup": 2.6652, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_a2f8_d352_smallk_boundary_b4_n1024_k256_d352", + "B": 4, + "N": 1024, + "D": 352, + "K": 256, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.176096, + "evolution_flashlib_ms": 0.479503, + "evolution_tflops": 4.192, + "evolution_flashlib_equiv_tflops": 1.5395, + "evolution_speedup": 2.723, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_8f09_d352_smallk_large_n_b4_n4096_k256_d352", + "B": 4, + "N": 4096, + "D": 352, + "K": 256, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.176128, + "evolution_flashlib_ms": 0.306655, + "evolution_tflops": 16.765, + "evolution_flashlib_equiv_tflops": 9.629, + "evolution_speedup": 1.7411, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d352_b4_n8192_k1024_d352", + "B": 4, + "N": 8192, + "D": 352, + "K": 1024, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.220992, + "evolution_flashlib_ms": 0.651839, + "evolution_tflops": 106.8922, + "evolution_flashlib_equiv_tflops": 36.2395, + "evolution_speedup": 2.9496, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_d9d5_d352_random_b5_n2304_k768_d352", + "B": 5, + "N": 2304, + "D": 352, + "K": 768, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.186944, + "evolution_flashlib_ms": 0.582592, + "evolution_tflops": 33.3177, + "evolution_flashlib_equiv_tflops": 10.6911, + "evolution_speedup": 3.1164, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d416_b1_n256_k256_d416", + "B": 1, + "N": 256, + "D": 416, + "K": 256, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.175137, + "evolution_flashlib_ms": 0.647807, + "evolution_tflops": 0.3113, + "evolution_flashlib_equiv_tflops": 0.0842, + "evolution_speedup": 3.6989, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_3328_d416_highk_low_n_b1_n384_k8192_d416", + "B": 1, + "N": 384, + "D": 416, + "K": 8192, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.276031, + "evolution_flashlib_ms": 41.538691, + "evolution_tflops": 9.4817, + "evolution_flashlib_equiv_tflops": 0.063, + "evolution_speedup": 150.4856, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d416_b1_n512_k8192_d416", + "B": 1, + "N": 512, + "D": 416, + "K": 8192, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.275103, + "evolution_flashlib_ms": 40.967689, + "evolution_tflops": 12.6849, + "evolution_flashlib_equiv_tflops": 0.0852, + "evolution_speedup": 148.9176, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_9ca0_d416_splitk_min_b1_n1024_k4096_d416", + "B": 1, + "N": 1024, + "D": 416, + "K": 4096, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.274271, + "evolution_flashlib_ms": 41.154087, + "evolution_tflops": 12.7234, + "evolution_flashlib_equiv_tflops": 0.0848, + "evolution_speedup": 150.049, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_ef00_d416_random_midk_b1_n2176_k768_d416", + "B": 1, + "N": 2176, + "D": 416, + "K": 768, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.191648, + "evolution_flashlib_ms": 0.529311, + "evolution_tflops": 7.255, + "evolution_flashlib_equiv_tflops": 2.6268, + "evolution_speedup": 2.7619, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_d9d5_d416_highk_request_b2_n640_k8192_d416", + "B": 2, + "N": 640, + "D": 416, + "K": 8192, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.298432, + "evolution_flashlib_ms": 3.254428, + "evolution_tflops": 29.2333, + "evolution_flashlib_equiv_tflops": 2.6807, + "evolution_speedup": 10.9051, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d416_b2_n2048_k1024_d416", + "B": 2, + "N": 2048, + "D": 416, + "K": 1024, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.193983, + "evolution_flashlib_ms": 0.642463, + "evolution_tflops": 17.9895, + "evolution_flashlib_equiv_tflops": 5.4317, + "evolution_speedup": 3.312, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_68cf_d416_overlap_b2_n2304_k1024_d416", + "B": 2, + "N": 2304, + "D": 416, + "K": 1024, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.196895, + "evolution_flashlib_ms": 0.613215, + "evolution_tflops": 19.9389, + "evolution_flashlib_equiv_tflops": 6.4021, + "evolution_speedup": 3.1144, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_a2f8_d416_random_b2_n2560_k768_d416", + "B": 2, + "N": 2560, + "D": 416, + "K": 768, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.189631, + "evolution_flashlib_ms": 0.492399, + "evolution_tflops": 17.2522, + "evolution_flashlib_equiv_tflops": 6.6441, + "evolution_speedup": 2.5966, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_c44f_d416_highk_request_b3_n512_k8192_d416", + "B": 3, + "N": 512, + "D": 416, + "K": 8192, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.303423, + "evolution_flashlib_ms": 3.822732, + "evolution_tflops": 34.5029, + "evolution_flashlib_equiv_tflops": 2.7386, + "evolution_speedup": 12.5987, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_8f09_d416_midk_tail_b3_n3456_k768_d416", + "B": 3, + "N": 3456, + "D": 416, + "K": 768, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.189088, + "evolution_flashlib_ms": 0.493791, + "evolution_tflops": 35.0361, + "evolution_flashlib_equiv_tflops": 13.4164, + "evolution_speedup": 2.6114, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_5600_d416_smallk_boundary_b4_n2048_k256_d416", + "B": 4, + "N": 2048, + "D": 416, + "K": 256, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.175743, + "evolution_flashlib_ms": 0.370128, + "evolution_tflops": 9.9283, + "evolution_flashlib_equiv_tflops": 4.7141, + "evolution_speedup": 2.1061, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_1d49_d416_random_b4_n3840_k512_d416", + "B": 4, + "N": 3840, + "D": 416, + "K": 512, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.184, + "evolution_flashlib_ms": 0.385919, + "evolution_tflops": 35.5604, + "evolution_flashlib_equiv_tflops": 16.9546, + "evolution_speedup": 2.0974, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d416_b4_n8192_k1024_d416", + "B": 4, + "N": 8192, + "D": 416, + "K": 1024, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.230463, + "evolution_flashlib_ms": 0.710015, + "evolution_tflops": 121.1354, + "evolution_flashlib_equiv_tflops": 39.3193, + "evolution_speedup": 3.0808, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_3328_d480_min_boundary_b1_n128_k256_d480", + "B": 1, + "N": 128, + "D": 480, + "K": 256, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.177056, + "evolution_flashlib_ms": 0.488575, + "evolution_tflops": 0.1777, + "evolution_flashlib_equiv_tflops": 0.0644, + "evolution_speedup": 2.7594, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d480_b1_n256_k256_d480", + "B": 1, + "N": 256, + "D": 480, + "K": 256, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.176655, + "evolution_flashlib_ms": 0.649151, + "evolution_tflops": 0.3561, + "evolution_flashlib_equiv_tflops": 0.0969, + "evolution_speedup": 3.6747, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d480_b1_n512_k8192_d480", + "B": 1, + "N": 512, + "D": 480, + "K": 8192, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.278559, + "evolution_flashlib_ms": 32.836689, + "evolution_tflops": 14.4548, + "evolution_flashlib_equiv_tflops": 0.1226, + "evolution_speedup": 117.8803, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_a2f8_d480_highk_request_b1_n896_k4096_d480", + "B": 1, + "N": 896, + "D": 480, + "K": 4096, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.274432, + "evolution_flashlib_ms": 32.096809, + "evolution_tflops": 12.8382, + "evolution_flashlib_equiv_tflops": 0.1098, + "evolution_speedup": 116.9572, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_5600_d480_random_b1_n1536_k1024_d480", + "B": 1, + "N": 1536, + "D": 480, + "K": 1024, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.199488, + "evolution_flashlib_ms": 0.718975, + "evolution_tflops": 7.5691, + "evolution_flashlib_equiv_tflops": 2.1001, + "evolution_speedup": 3.6041, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_8f09_d480_highk_low_n_b2_n640_k4096_d480", + "B": 2, + "N": 640, + "D": 480, + "K": 4096, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.284192, + "evolution_flashlib_ms": 1.999325, + "evolution_tflops": 17.7104, + "evolution_flashlib_equiv_tflops": 2.5174, + "evolution_speedup": 7.0351, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d480_b2_n2048_k1024_d480", + "B": 2, + "N": 2048, + "D": 480, + "K": 1024, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.197919, + "evolution_flashlib_ms": 0.672447, + "evolution_tflops": 20.3443, + "evolution_flashlib_equiv_tflops": 5.9879, + "evolution_speedup": 3.3976, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_1d49_d480_smallk_boundary_b2_n2816_k256_d480", + "B": 2, + "N": 2816, + "D": 480, + "K": 256, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.178015, + "evolution_flashlib_ms": 0.332511, + "evolution_tflops": 7.7753, + "evolution_flashlib_equiv_tflops": 4.1626, + "evolution_speedup": 1.8679, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_9ca0_d480_splitd_large_n_b2_n4096_k512_d480", + "B": 2, + "N": 4096, + "D": 480, + "K": 512, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.185087, + "evolution_flashlib_ms": 0.447807, + "evolution_tflops": 21.7548, + "evolution_flashlib_equiv_tflops": 8.9917, + "evolution_speedup": 2.4194, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_d9d5_d480_heldout_highk_b3_n1024_k4096_d480", + "B": 3, + "N": 1024, + "D": 480, + "K": 4096, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.297632, + "evolution_flashlib_ms": 1.92259, + "evolution_tflops": 40.5857, + "evolution_flashlib_equiv_tflops": 6.283, + "evolution_speedup": 6.4596, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_ef00_d480_lowk_boundary_b3_n3200_k256_d480", + "B": 3, + "N": 3200, + "D": 480, + "K": 256, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.18192, + "evolution_flashlib_ms": 0.334224, + "evolution_tflops": 12.9689, + "evolution_flashlib_equiv_tflops": 7.059, + "evolution_speedup": 1.8372, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_68cf_d480_boundary_b4_n1664_k512_d480", + "B": 4, + "N": 1664, + "D": 480, + "K": 512, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.18536, + "evolution_flashlib_ms": 0.539839, + "evolution_tflops": 17.6497, + "evolution_flashlib_equiv_tflops": 6.0602, + "evolution_speedup": 2.9124, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "post_d895_d480_b4_n8192_k1024_d480", + "B": 4, + "N": 8192, + "D": 480, + "K": 1024, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.24, + "evolution_flashlib_ms": 0.733983, + "evolution_tflops": 134.2177, + "evolution_flashlib_equiv_tflops": 43.8869, + "evolution_speedup": 3.0583, + }, + { + "source": "post_d895_plus_ledgers_124", + "label": "adjacent_c44f_d480_boundary_b5_n2048_k512_d480", + "B": 5, + "N": 2048, + "D": 480, + "K": 512, + "dtype": "bfloat16", + "route": "gap_pad_to_supported_seed_v1", + "evolution_kernel_ms": 0.185136, + "evolution_flashlib_ms": 0.501759, + "evolution_tflops": 27.1863, + "evolution_flashlib_equiv_tflops": 10.031, + "evolution_speedup": 2.7102, + }, +] diff --git a/cake_exports/kmeans/benchmarks/flash_kmeans_triton_h200.py b/cake_exports/kmeans/benchmarks/flash_kmeans_triton_h200.py new file mode 100644 index 00000000..72cb364a --- /dev/null +++ b/cake_exports/kmeans/benchmarks/flash_kmeans_triton_h200.py @@ -0,0 +1,227 @@ +from __future__ import annotations + +from typing import Any + +import torch +import triton +import triton.language as tl + + +def _next_pow2(v: int) -> int: + if v <= 1: + return 1 + return 1 << (v - 1).bit_length() + + +def _pad_d(dim: int) -> int: + return max(16, _next_pow2(int(dim))) + + +def h200_small_d_config(n_points: int, n_clusters: int, dim: int, dtype: Any) -> dict[str, int]: + """Flash-KMeans H200 fp16/bf16 small-D heuristic. + + This is copied from flash-kmeans' H200 heuristic and intentionally ignores + the local GPU name. The benchmark uses it as the default Triton baseline + even on B200/B300 so the comparison is stable and matches the MR audit. + """ + + if dtype not in {torch.float16, torch.bfloat16}: + raise ValueError(f"H200 baseline in this benchmark expects fp16/bf16, got {dtype}") + if dim > 512: + raise ValueError(f"H200 small-D Triton baseline supports D <= 512, got D={dim}") + + block_n = 128 + block_k = 64 + num_warps = 4 + num_stages = 1 + + if dim >= 512: + block_n = 128 + block_k = 64 + num_warps = 8 + num_stages = 1 + elif dim >= 256: + block_n = 128 + block_k = 64 + num_warps = 4 + num_stages = 2 + else: + if n_clusters >= 4096: + block_k = 128 + if dim >= 128: + num_warps = 8 + num_stages = 2 + else: + num_warps = 4 + num_stages = 4 + else: + block_k = 64 + num_warps = 4 + num_stages = 1 + + if dim <= 64 and n_clusters >= 4096: + block_n = 64 + block_k = 128 + num_warps = 4 + num_stages = 4 + + if n_points < 65536: + block_n = 64 + + return { + "kernel": "triton_h200_small_d", + "BLOCK_N": int(block_n), + "BLOCK_K": int(block_k), + "D_PAD": int(_pad_d(dim)), + "num_warps": int(num_warps), + "num_stages": int(num_stages), + } + + +@triton.jit +def _euclid_assign_kernel( + x_ptr, + c_ptr, + x_sq_ptr, + c_sq_ptr, + out_ptr, + B: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, + D: tl.constexpr, + stride_x_b: tl.constexpr, + stride_x_n: tl.constexpr, + stride_x_d: tl.constexpr, + stride_c_b: tl.constexpr, + stride_c_k: tl.constexpr, + stride_c_d: tl.constexpr, + stride_xsq_b: tl.constexpr, + stride_xsq_n: tl.constexpr, + stride_csq_b: tl.constexpr, + stride_csq_k: tl.constexpr, + stride_out_b: tl.constexpr, + stride_out_n: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + D_PAD: tl.constexpr, +): + pid_n = tl.program_id(0) + pid_b = tl.program_id(1).to(tl.int64) + + n_start = pid_n * BLOCK_N + n_offsets = (n_start + tl.arange(0, BLOCK_N)).to(tl.int64) + n_mask = n_offsets < N + offs_d = tl.arange(0, D_PAD).to(tl.int64) + d_mask = offs_d < D + + x_ptrs = ( + x_ptr + + pid_b * stride_x_b + + n_offsets[:, None] * stride_x_n + + offs_d[None, :] * stride_x_d + ) + x_tile = tl.load(x_ptrs, mask=n_mask[:, None] & d_mask[None, :], other=0.0) + + xsq_ptrs = x_sq_ptr + pid_b * stride_xsq_b + n_offsets * stride_xsq_n + x_sq_tile = tl.load(xsq_ptrs, mask=n_mask, other=0.0).to(tl.float32) + + best_dist = tl.full((BLOCK_N,), 3.4e38, tl.float32) + best_idx = tl.zeros((BLOCK_N,), tl.int32) + + for k_start in range(0, K, BLOCK_K): + k_offsets = (k_start + tl.arange(0, BLOCK_K)).to(tl.int64) + k_mask = k_offsets < K + + c_ptrs = ( + c_ptr + + pid_b * stride_c_b + + k_offsets[None, :] * stride_c_k + + offs_d[:, None] * stride_c_d + ) + c_tile = tl.load(c_ptrs, mask=k_mask[None, :] & d_mask[:, None], other=0.0) + + csq_ptrs = c_sq_ptr + pid_b * stride_csq_b + k_offsets * stride_csq_k + c_sq_tile = tl.load(csq_ptrs, mask=k_mask, other=0.0).to(tl.float32) + + cross = tl.dot(x_tile, c_tile).to(tl.float32) + dist = x_sq_tile[:, None] + c_sq_tile[None, :] - 2.0 * cross + dist = tl.maximum(dist, 0.0) + dist = tl.where(k_mask[None, :], dist, 3.4e38) + + curr_min = tl.min(dist, axis=1) + curr_idx = tl.argmin(dist, axis=1) + update = curr_min < best_dist + best_dist = tl.where(update, curr_min, best_dist) + best_idx = tl.where(update, k_start + curr_idx, best_idx) + + out_ptrs = out_ptr + pid_b * stride_out_b + n_offsets * stride_out_n + tl.store(out_ptrs, best_idx, mask=n_mask) + + +def euclid_assign_triton_h200( + x: torch.Tensor, + centroids: torch.Tensor, + x_sq: torch.Tensor, + c_sq: torch.Tensor, + *, + out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, dict[str, int]]: + """Run Flash-KMeans' Triton Euclidean assign with forced H200 heuristic.""" + + if x.ndim != 3 or centroids.ndim != 3: + raise ValueError("x and centroids must have shapes [B, N, D] and [B, K, D]") + if x.dtype not in {torch.float16, torch.bfloat16}: + raise ValueError(f"Triton H200 baseline expects fp16/bf16 inputs, got {x.dtype}") + if x.device != centroids.device: + raise ValueError("x and centroids must be on the same device") + if not x.is_cuda or not centroids.is_cuda: + raise ValueError("x and centroids must be CUDA tensors") + if not x.is_contiguous() or not centroids.is_contiguous(): + raise ValueError("x and centroids must be contiguous") + + bsz, n_points, dim = (int(v) for v in x.shape) + n_clusters = int(centroids.shape[1]) + if int(centroids.shape[0]) != bsz or int(centroids.shape[2]) != dim: + raise ValueError("x and centroids must have matching B and D dimensions") + if out is None: + out = torch.empty((bsz, n_points), dtype=torch.int32, device=x.device) + if tuple(out.shape) != (bsz, n_points) or out.dtype is not torch.int32: + raise ValueError("out must have shape [B, N] and dtype torch.int32") + + cfg = h200_small_d_config(n_points, n_clusters, dim, x.dtype) + stride_x_b, stride_x_n, stride_x_d = x.stride() + stride_c_b, stride_c_k, stride_c_d = centroids.stride() + stride_xsq_b, stride_xsq_n = x_sq.stride() + stride_csq_b, stride_csq_k = c_sq.stride() + stride_out_b, stride_out_n = out.stride() + + grid = lambda meta: (triton.cdiv(n_points, meta["BLOCK_N"]), bsz) + _euclid_assign_kernel[grid]( + x, + centroids, + x_sq, + c_sq, + out, + bsz, + n_points, + n_clusters, + dim, + stride_x_b, + stride_x_n, + stride_x_d, + stride_c_b, + stride_c_k, + stride_c_d, + stride_xsq_b, + stride_xsq_n, + stride_csq_b, + stride_csq_k, + stride_out_b, + stride_out_n, + BLOCK_N=cfg["BLOCK_N"], + BLOCK_K=cfg["BLOCK_K"], + D_PAD=cfg["D_PAD"], + num_warps=cfg["num_warps"], + num_stages=cfg["num_stages"], + ) + return out, cfg diff --git a/cake_exports/kmeans/benchmarks/flash_kmeans_triton_h200_raw_adapter.py b/cake_exports/kmeans/benchmarks/flash_kmeans_triton_h200_raw_adapter.py new file mode 100644 index 00000000..b78efbc9 --- /dev/null +++ b/cake_exports/kmeans/benchmarks/flash_kmeans_triton_h200_raw_adapter.py @@ -0,0 +1,251 @@ +"""Fair raw-input runtime adapter around the frozen 07cf Triton assignment.""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass, field +from threading import Condition, RLock +from typing import Any + +from flash_kmeans_triton_h200 import euclid_assign_triton_h200 + +BASELINE_NAME = "triton_h200_07cf_raw_adapter_v1" +BASELINE_COMMIT = "07cf2a27928aacf6790c950a265d8b8dc83c87cf" +PREPROCESS_IMPL = "flashlib_cake_kmeans._row_norm:PreparedBF16PairRowNorm" + + +def _record_stream(stream: Any, *tensors: Any) -> None: + """Keep caller and scratch storage allocator-safe after async submission.""" + + seen: set[int] = set() + for tensor in tensors: + identity = id(tensor) + if identity in seen: + continue + seen.add(identity) + record_stream = getattr(tensor, "record_stream", None) + if callable(record_stream): + record_stream(stream) + + +@dataclass +class _Slot: + x_sq: Any + c_sq: Any + norm_plan: Any + lock: RLock = field(default_factory=RLock, repr=False) + + +class TritonH20007cfRawAdapter: + """Init-once raw-input adapter around the frozen 07cf assignment kernel. + + The adapter intentionally reuses the candidate's exported pair-row-norm + implementation. The frozen 07cf assignment consumes both norm fields, + while a candidate route may elide fields that it does not bind; this lane + therefore compares complete raw-input operators, not assignment-only + stages. Scratch and prepared norm launches are isolated by shape and CUDA + stream. Each public call allocates its output before preprocessing, matching + ``FlashKMeansAssignRuntime.compute`` ordering and return semantics. + """ + + def __init__(self, *, device_index: int, arch: str) -> None: + self.device_index = int(device_index) + self.arch = str(arch) + self._slots: dict[tuple[Any, ...], _Slot] = {} + self._cache_lock = RLock() + self._lifecycle = Condition(RLock()) + self._active_computes = 0 + self._clearing = False + self._hits = 0 + self._misses = 0 + + @contextmanager + def _compute_lifecycle(self): + with self._lifecycle: + while self._clearing: + self._lifecycle.wait() + self._active_computes += 1 + try: + yield + finally: + with self._lifecycle: + self._active_computes -= 1 + if self._active_computes == 0: + self._lifecycle.notify_all() + + def cache_info(self) -> dict[str, int]: + with self._cache_lock: + return { + "size": len(self._slots), + "hits": self._hits, + "misses": self._misses, + } + + def clear(self, *, synchronize: bool = True) -> None: + """Exclusively release scratch after admitted submissions finish.""" + + import torch + + with self._lifecycle: + while self._clearing: + self._lifecycle.wait() + try: + self._clearing = True + while self._active_computes: + self._lifecycle.wait() + if synchronize: + with torch.cuda.device(self.device_index): + torch.cuda.synchronize() + with self._cache_lock: + self._slots.clear() + self._hits = 0 + self._misses = 0 + finally: + self._clearing = False + self._lifecycle.notify_all() + + def _resolve_stream(self, x: Any, centroids: Any, stream: Any) -> tuple[Any, tuple[Any, ...]]: + import torch + + if not all(isinstance(item, torch.Tensor) and item.is_cuda for item in (x, centroids)): + raise TypeError("07cf raw adapter inputs must be CUDA torch.Tensor objects") + if x.dtype is not torch.bfloat16 or centroids.dtype is not torch.bfloat16: + raise TypeError("07cf raw adapter inputs must have bfloat16 dtype") + if x.ndim != 3 or centroids.ndim != 3 or not x.is_contiguous() or not centroids.is_contiguous(): + raise ValueError("07cf raw adapter inputs must be contiguous [B, rows, D] tensors") + bsz, n_points, dim = map(int, x.shape) + c_bsz, n_clusters, c_dim = map(int, centroids.shape) + if (bsz, dim) != (c_bsz, c_dim) or x.device != centroids.device: + raise ValueError("07cf raw adapter inputs must have matching batch, feature, and device") + input_device_index = x.device.index + if input_device_index is None: + input_device_index = torch.cuda.current_device() + input_device_index = int(input_device_index) + if input_device_index != self.device_index: + raise ValueError( + f"07cf raw adapter targets CUDA device {self.device_index}, got input device {input_device_index}" + ) + with torch.cuda.device(self.device_index): + resolved_stream = torch.cuda.current_stream(self.device_index) if stream is None else stream + stream_device = getattr(resolved_stream, "device", None) + stream_device_index = getattr(stream_device, "index", stream_device) + if stream_device_index is not None and int(stream_device_index) != self.device_index: + raise ValueError( + f"07cf raw adapter stream device {stream_device_index} does not match input device {self.device_index}" + ) + stream_handle = int(resolved_stream.cuda_stream) + key = ( + self.device_index, + self.arch, + bsz, + n_points, + n_clusters, + dim, + str(x.dtype), + stream_handle, + ) + return resolved_stream, key + + def _locked_slot( + self, + key: tuple[Any, ...], + x: Any, + centroids: Any, + *, + stream: Any, + ) -> tuple[_Slot, bool]: + import torch + from flashlib_cake_kmeans._row_norm import prepare_bf16_pair_row_norm + + # Miss preparation is serialized under the cache lock. This is a cold + # path and prevents two threads from compiling/publishing duplicate + # launch plans for the same shape/stream key. + with self._cache_lock: + slot = self._slots.get(key) + cache_hit = slot is not None + if slot is None: + with torch.cuda.device(self.device_index), torch.cuda.stream(stream): + bsz, n_points, _ = map(int, x.shape) + n_clusters = int(centroids.shape[1]) + x_sq = torch.empty((bsz, n_points), dtype=torch.float32, device=x.device) + c_sq = torch.empty((bsz, n_clusters), dtype=torch.float32, device=x.device) + norm_plan = prepare_bf16_pair_row_norm( + x, + centroids, + x_sq, + c_sq, + compute_x=True, + compute_c=True, + arch=self.arch, + stream=stream, + ) + slot = _Slot(x_sq=x_sq, c_sq=c_sq, norm_plan=norm_plan) + self._slots[key] = slot + self._misses += 1 + else: + self._hits += 1 + slot.lock.acquire() + return slot, cache_hit + + def compute( + self, + x: Any, + centroids: Any, + *, + stream: Any = None, + return_info: bool = False, + ): + import torch + + with self._compute_lifecycle(): + resolved_stream, key = self._resolve_stream(x, centroids, stream) + with torch.cuda.device(self.device_index), torch.cuda.stream(resolved_stream): + bsz, n_points, _ = map(int, x.shape) + out = torch.empty((bsz, n_points), dtype=torch.int32, device=x.device) + slot, cache_hit = self._locked_slot( + key, + x, + centroids, + stream=resolved_stream, + ) + try: + with torch.cuda.device(self.device_index), torch.cuda.stream(resolved_stream): + _record_stream(resolved_stream, x, centroids, slot.x_sq, slot.c_sq, out) + try: + slot.norm_plan.rebind( + x, + centroids, + slot.x_sq, + slot.c_sq, + stream=resolved_stream, + ) + slot.norm_plan.launch(stream=resolved_stream) + finally: + # The prepared driver launch must not keep caller-owned + # x/centroids alive after enqueue returns. + slot.norm_plan.release_bound_callers( + slot.x_sq, + stream=resolved_stream, + ) + produced_out, config = euclid_assign_triton_h200( + x, + centroids, + slot.x_sq, + slot.c_sq, + out=out, + ) + if produced_out is not out: + raise RuntimeError("frozen 07cf assignment did not preserve the adapter-owned output") + finally: + slot.lock.release() + if not return_info: + return out + return out, { + "triton_h200_07cf_config": config, + "runtime_cache_hit": cache_hit, + "norm_launch_count": 1, + "norm_compute_fields": ("x_sq", "c_sq"), + "assignment_launch_count": 1, + "runtime_launch_count": 2, + "stream_handle": int(resolved_stream.cuda_stream), + } diff --git a/cake_exports/kmeans/benchmarks/shape_records.json b/cake_exports/kmeans/benchmarks/shape_records.json new file mode 100644 index 00000000..cd982b94 --- /dev/null +++ b/cake_exports/kmeans/benchmarks/shape_records.json @@ -0,0 +1,4458 @@ +[ + { + "label": "d895_expanded_heldout_neighborhood_d80_b2_n2176_k1024_d80", + "params": { + "B": 2, + "D": 80, + "K": 1024, + "N": 2176, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.7696, + "evolution_flashlib_ms": 0.402944, + "evolution_kernel_ms": 0.176928, + "evolution_speedup": 2.2774, + "evolution_tflops": 4.0301, + "route": "lowdim_e50c_v1", + "source": "expanded9" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d895_expanded_guard_miss_fallback_d128_b1_n1408_k512_d128", + "params": { + "B": 1, + "D": 128, + "K": 512, + "N": 1408, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.53, + "evolution_flashlib_ms": 0.348175, + "evolution_kernel_ms": 0.144544, + "evolution_speedup": 2.4088, + "evolution_tflops": 1.2768, + "route": "aligned_v10_fallback", + "source": "expanded9" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d895_expanded_forced_fallback_d128_b1_n1664_k256_d128", + "params": { + "B": 1, + "D": 128, + "K": 256, + "N": 1664, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.3924, + "evolution_flashlib_ms": 0.277904, + "evolution_kernel_ms": 0.142032, + "evolution_speedup": 1.9566, + "evolution_tflops": 0.7678, + "route": "aligned_v10_fallback", + "source": "expanded9" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d895_expanded_guard_boundary_d144_b1_n128_k256_d144", + "params": { + "B": 1, + "D": 144, + "K": 256, + "N": 128, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0205, + "evolution_flashlib_ms": 0.459647, + "evolution_kernel_ms": 0.1664, + "evolution_speedup": 2.7623, + "evolution_tflops": 0.0567, + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "source": "expanded9" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d895_expanded_tail_divisibility_d176_b1_n1152_k512_d176", + "params": { + "B": 1, + "D": 176, + "K": 512, + "N": 1152, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.5253, + "evolution_flashlib_ms": 0.395232, + "evolution_kernel_ms": 0.174944, + "evolution_speedup": 2.2592, + "evolution_tflops": 1.1868, + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "source": "expanded9" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d895_expanded_guard_overlap_d192_b1_n1024_k768_d192", + "params": { + "B": 1, + "D": 192, + "K": 768, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.6618, + "evolution_flashlib_ms": 0.45632, + "evolution_kernel_ms": 0.17168, + "evolution_speedup": 2.658, + "evolution_tflops": 1.759, + "route": "d192_paired_repeated_mma_v1", + "source": "expanded9" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d895_expanded_random_legal_d320_b1_n1280_k512_d320", + "params": { + "B": 1, + "D": 320, + "K": 512, + "N": 1280, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.826, + "evolution_flashlib_ms": 0.507759, + "evolution_kernel_ms": 0.152496, + "evolution_speedup": 3.3297, + "evolution_tflops": 2.7504, + "route": "highd_splitd_single_tile_6fcf_v1", + "source": "expanded9" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d895_expanded_request_specific_d448_b1_n1280_k1024_d448", + "params": { + "B": 1, + "D": 448, + "K": 1024, + "N": 1280, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.7411, + "evolution_flashlib_ms": 0.674512, + "evolution_kernel_ms": 0.173743, + "evolution_speedup": 3.8822, + "evolution_tflops": 6.7594, + "route": "highd_splitd_single_tile_6fcf_v1", + "source": "expanded9" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d895_expanded_request_specific_d512_splitk_b1_n512_k8192_d512", + "params": { + "B": 1, + "D": 512, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1375, + "evolution_flashlib_ms": 31.230383, + "evolution_kernel_ms": 0.257392, + "evolution_speedup": 121.3339, + "evolution_tflops": 16.6865, + "route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "source": "expanded9" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d16_paired_b2_n4096_k1024_d16", + "params": { + "B": 2, + "D": 16, + "K": 1024, + "N": 4096, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.8529, + "evolution_flashlib_ms": 0.314751, + "evolution_kernel_ms": 0.175568, + "evolution_speedup": 1.7928, + "evolution_tflops": 1.529, + "route": "microdim_hybrid_9c0d_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d16_fallback_b3_n2432_k512_d16", + "params": { + "B": 3, + "D": 16, + "K": 512, + "N": 2432, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.4098, + "evolution_flashlib_ms": 0.291712, + "evolution_kernel_ms": 0.127456, + "evolution_speedup": 2.2887, + "evolution_tflops": 0.9379, + "route": "microdim_hybrid_9c0d_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d16_small_b4_n1024_k512_d16", + "params": { + "B": 4, + "D": 16, + "K": 512, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1836, + "evolution_flashlib_ms": 0.365599, + "evolution_kernel_ms": 0.141568, + "evolution_speedup": 2.5825, + "evolution_tflops": 0.474, + "route": "microdim_hybrid_9c0d_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d16_large_b8_n32768_k512_d16", + "params": { + "B": 8, + "D": 16, + "K": 512, + "N": 32768, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 13.0308, + "evolution_flashlib_ms": 0.3296, + "evolution_kernel_ms": 0.234416, + "evolution_speedup": 1.406, + "evolution_tflops": 18.322, + "route": "microdim_hybrid_9c0d_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d32_hugek_b1_n512_k8192_d32", + "params": { + "B": 1, + "D": 32, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1923, + "evolution_flashlib_ms": 1.396063, + "evolution_kernel_ms": 0.224064, + "evolution_speedup": 6.2306, + "evolution_tflops": 1.198, + "route": "microdim_hybrid_9c0d_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d32_paired_b2_n4096_k1024_d32", + "params": { + "B": 2, + "D": 32, + "K": 1024, + "N": 4096, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.6899, + "evolution_flashlib_ms": 0.317696, + "evolution_kernel_ms": 0.176672, + "evolution_speedup": 1.7982, + "evolution_tflops": 3.0388, + "route": "microdim_hybrid_9c0d_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d32_fallback_b3_n2432_k512_d32", + "params": { + "B": 3, + "D": 32, + "K": 512, + "N": 2432, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.8475, + "evolution_flashlib_ms": 0.28208, + "evolution_kernel_ms": 0.133824, + "evolution_speedup": 2.1078, + "evolution_tflops": 1.7865, + "route": "microdim_hybrid_9c0d_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d32_small_b4_n1024_k512_d32", + "params": { + "B": 4, + "D": 32, + "K": 512, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.3716, + "evolution_flashlib_ms": 0.361167, + "evolution_kernel_ms": 0.139583, + "evolution_speedup": 2.5875, + "evolution_tflops": 0.9616, + "route": "microdim_hybrid_9c0d_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d64_paired_b2_n4096_k1024_d64", + "params": { + "B": 2, + "D": 64, + "K": 1024, + "N": 4096, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 3.3518, + "evolution_flashlib_ms": 0.320352, + "evolution_kernel_ms": 0.070272, + "evolution_speedup": 4.5587, + "evolution_tflops": 15.2798, + "route": "d64_direct_single64_1p2gap_9f2a_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d64_fallback_b3_n2432_k512_d64", + "params": { + "B": 3, + "D": 64, + "K": 512, + "N": 2432, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.6742, + "evolution_flashlib_ms": 0.2856, + "evolution_kernel_ms": 0.067552, + "evolution_speedup": 4.2279, + "evolution_tflops": 7.0783, + "route": "d64_direct_single64_1p2gap_9f2a_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d64_small_b4_n1024_k512_d64", + "params": { + "B": 4, + "D": 64, + "K": 512, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.6741, + "evolution_flashlib_ms": 0.398223, + "evolution_kernel_ms": 0.06736, + "evolution_speedup": 5.9119, + "evolution_tflops": 3.9851, + "route": "d64_direct_single64_1p2gap_9f2a_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d64_large_b8_n32768_k512_d64", + "params": { + "B": 8, + "D": 64, + "K": 512, + "N": 32768, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 34.8042, + "evolution_flashlib_ms": 0.493615, + "evolution_kernel_ms": 0.128672, + "evolution_speedup": 3.8362, + "evolution_tflops": 133.5168, + "route": "d64_direct_single64_1p2gap_9f2a_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d80_paired_b1_n2048_k1024_d80", + "params": { + "B": 1, + "D": 80, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.8736, + "evolution_flashlib_ms": 0.384112, + "evolution_kernel_ms": 0.176896, + "evolution_speedup": 2.1714, + "evolution_tflops": 1.8968, + "route": "lowdim_e50c_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d80_small_b2_n1024_k512_d80", + "params": { + "B": 2, + "D": 80, + "K": 512, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.4431, + "evolution_flashlib_ms": 0.378656, + "evolution_kernel_ms": 0.173552, + "evolution_speedup": 2.1818, + "evolution_tflops": 0.9667, + "route": "lowdim_e50c_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d96_paired_b1_n1536_k1024_d96", + "params": { + "B": 1, + "D": 96, + "K": 1024, + "N": 1536, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.7061, + "evolution_flashlib_ms": 0.427679, + "evolution_kernel_ms": 0.176032, + "evolution_speedup": 2.4296, + "evolution_tflops": 1.7155, + "route": "lowdim_e50c_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d96_small_b2_n896_k512_d96", + "params": { + "B": 2, + "D": 96, + "K": 512, + "N": 896, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.4465, + "evolution_flashlib_ms": 0.394576, + "evolution_kernel_ms": 0.173504, + "evolution_speedup": 2.2742, + "evolution_tflops": 1.0153, + "route": "lowdim_e50c_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d96_random_b3_n6144_k1280_d96", + "params": { + "B": 3, + "D": 96, + "K": 1280, + "N": 6144, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 10.6231, + "evolution_flashlib_ms": 0.426416, + "evolution_kernel_ms": 0.177823, + "evolution_speedup": 2.398, + "evolution_tflops": 25.4739, + "route": "lowdim_e50c_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d96_fallback_b5_n1664_k256_d96", + "params": { + "B": 5, + "D": 96, + "K": 256, + "N": 1664, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.3432, + "evolution_flashlib_ms": 0.304464, + "evolution_kernel_ms": 0.171136, + "evolution_speedup": 1.7791, + "evolution_tflops": 2.3896, + "route": "lowdim_e50c_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "boundary_b1_n128_k256_d128", + "params": { + "B": 1, + "D": 128, + "K": 256, + "N": 128, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0218, + "evolution_flashlib_ms": 0.384415, + "evolution_kernel_ms": 0.143039, + "evolution_speedup": 2.6875, + "evolution_tflops": 0.0586, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_tiny_hugek_b1_n128_k4096_d128", + "params": { + "B": 1, + "D": 128, + "K": 4096, + "N": 128, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1114, + "evolution_flashlib_ms": 1.205311, + "evolution_kernel_ms": 0.171472, + "evolution_speedup": 7.0292, + "evolution_tflops": 0.7827, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "small_b1_n256_k256_d128", + "params": { + "B": 1, + "D": 128, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0385, + "evolution_flashlib_ms": 0.435263, + "evolution_kernel_ms": 0.145312, + "evolution_speedup": 2.9954, + "evolution_tflops": 0.1155, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_min_even_b1_n256_k768_d128", + "params": { + "B": 1, + "D": 128, + "K": 768, + "N": 256, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.084, + "evolution_flashlib_ms": 0.599455, + "evolution_kernel_ms": 0.167616, + "evolution_speedup": 3.5764, + "evolution_tflops": 0.3003, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_low_n_hugek_b1_n256_k4096_d128", + "params": { + "B": 1, + "D": 128, + "K": 4096, + "N": 256, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1873, + "evolution_flashlib_ms": 1.432831, + "evolution_kernel_ms": 0.209535, + "evolution_speedup": 6.8381, + "evolution_tflops": 1.2811, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "small_b1_n896_k512_d128", + "params": { + "B": 1, + "D": 128, + "K": 512, + "N": 896, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.3102, + "evolution_flashlib_ms": 0.378655, + "evolution_kernel_ms": 0.144463, + "evolution_speedup": 2.6211, + "evolution_tflops": 0.8129, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_odd_kover_b1_n896_k768_d128", + "params": { + "B": 1, + "D": 128, + "K": 768, + "N": 896, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.3949, + "evolution_flashlib_ms": 0.446111, + "evolution_kernel_ms": 0.145856, + "evolution_speedup": 3.0586, + "evolution_tflops": 1.2078, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_hugek_b1_n896_k16384_d128", + "params": { + "B": 1, + "D": 128, + "K": 16384, + "N": 896, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.6987, + "evolution_flashlib_ms": 2.21235, + "evolution_kernel_ms": 0.25696, + "evolution_speedup": 8.6097, + "evolution_tflops": 14.6252, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "boundary_b1_n1024_k768_d128", + "params": { + "B": 1, + "D": 128, + "K": 768, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.4318, + "evolution_flashlib_ms": 0.466272, + "evolution_kernel_ms": 0.169952, + "evolution_speedup": 2.7435, + "evolution_tflops": 1.1846, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_kover_b1_n1024_k2048_d128", + "params": { + "B": 1, + "D": 128, + "K": 2048, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.8902, + "evolution_flashlib_ms": 0.603104, + "evolution_kernel_ms": 0.1848, + "evolution_speedup": 3.2635, + "evolution_tflops": 2.9051, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "forced_fallback_b1_n1152_k256_d128", + "params": { + "B": 1, + "D": 128, + "K": 256, + "N": 1152, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.2411, + "evolution_flashlib_ms": 0.313135, + "evolution_kernel_ms": 0.143648, + "evolution_speedup": 2.1799, + "evolution_tflops": 0.5256, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_first_after_small_b1_n1152_k512_d128", + "params": { + "B": 1, + "D": 128, + "K": 512, + "N": 1152, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.3964, + "evolution_flashlib_ms": 0.38096, + "evolution_kernel_ms": 0.144832, + "evolution_speedup": 2.6304, + "evolution_tflops": 1.0426, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "boundary_b1_n1280_k256_d128", + "params": { + "B": 1, + "D": 128, + "K": 256, + "N": 1280, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.2722, + "evolution_flashlib_ms": 0.308176, + "evolution_kernel_ms": 0.163984, + "evolution_speedup": 1.8793, + "evolution_tflops": 0.5116, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_even_b1_n1536_k256_d128", + "params": { + "B": 1, + "D": 128, + "K": 256, + "N": 1536, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.3641, + "evolution_flashlib_ms": 0.276479, + "evolution_kernel_ms": 0.160624, + "evolution_speedup": 1.7213, + "evolution_tflops": 0.6267, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_largek_b1_n2048_k4096_d128", + "params": { + "B": 1, + "D": 128, + "K": 4096, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.9421, + "evolution_flashlib_ms": 0.729904, + "evolution_kernel_ms": 0.209024, + "evolution_speedup": 3.492, + "evolution_tflops": 10.2739, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_odd_b1_n2176_k4096_d128", + "params": { + "B": 1, + "D": 128, + "K": 4096, + "N": 2176, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 3.1992, + "evolution_flashlib_ms": 0.713215, + "evolution_kernel_ms": 0.169311, + "evolution_speedup": 4.2125, + "evolution_tflops": 13.4764, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "tail_odd_b1_n4224_k1024_d128", + "params": { + "B": 1, + "D": 128, + "K": 1024, + "N": 4224, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.9054, + "evolution_flashlib_ms": 0.38112, + "evolution_kernel_ms": 0.14784, + "evolution_speedup": 2.5779, + "evolution_tflops": 7.4898, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_large_odd_b1_n8320_k1280_d128", + "params": { + "B": 1, + "D": 128, + "K": 1280, + "N": 8320, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 7.8828, + "evolution_flashlib_ms": 0.345856, + "evolution_kernel_ms": 0.149984, + "evolution_speedup": 2.306, + "evolution_tflops": 18.1773, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_large_b1_n65536_k4096_d128", + "params": { + "B": 1, + "D": 128, + "K": 4096, + "N": 65536, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 62.6692, + "evolution_flashlib_ms": 1.096543, + "evolution_kernel_ms": 0.259808, + "evolution_speedup": 4.2206, + "evolution_tflops": 264.501, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_large_n_b1_n131072_k256_d128", + "params": { + "B": 1, + "D": 128, + "K": 256, + "N": 131072, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 31.6907, + "evolution_flashlib_ms": 0.271056, + "evolution_kernel_ms": 0.184608, + "evolution_speedup": 1.4683, + "evolution_tflops": 46.5307, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_kover_b2_n256_k768_d128", + "params": { + "B": 2, + "D": 128, + "K": 768, + "N": 256, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1597, + "evolution_flashlib_ms": 0.630431, + "evolution_kernel_ms": 0.16816, + "evolution_speedup": 3.749, + "evolution_tflops": 0.5986, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_low_n_hugek_b2_n384_k8192_d128", + "params": { + "B": 2, + "D": 128, + "K": 8192, + "N": 384, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.841, + "evolution_flashlib_ms": 1.915198, + "evolution_kernel_ms": 0.20048, + "evolution_speedup": 9.5531, + "evolution_tflops": 8.0338, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_even_b2_n512_k8192_d128", + "params": { + "B": 2, + "D": 128, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.2615, + "evolution_flashlib_ms": 1.702303, + "evolution_kernel_ms": 0.269792, + "evolution_speedup": 6.3097, + "evolution_tflops": 7.9598, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "small_random_b2_n768_k512_d128", + "params": { + "B": 2, + "D": 128, + "K": 512, + "N": 768, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.4722, + "evolution_flashlib_ms": 0.426368, + "evolution_kernel_ms": 0.144992, + "evolution_speedup": 2.9406, + "evolution_tflops": 1.3885, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "random_b2_n768_k3072_d128", + "params": { + "B": 2, + "D": 128, + "K": 3072, + "N": 768, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.4768, + "evolution_flashlib_ms": 0.817951, + "evolution_kernel_ms": 0.319424, + "evolution_speedup": 2.5607, + "evolution_tflops": 3.7817, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "smoke_b2_n1024_k512_d128", + "params": { + "B": 2, + "D": 128, + "K": 512, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.6361, + "evolution_flashlib_ms": 0.422031, + "evolution_kernel_ms": 0.148896, + "evolution_speedup": 2.8344, + "evolution_tflops": 1.8028, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_first_odd_b2_n1408_k256_d128", + "params": { + "B": 2, + "D": 128, + "K": 256, + "N": 1408, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.6573, + "evolution_flashlib_ms": 0.280752, + "evolution_kernel_ms": 0.14272, + "evolution_speedup": 1.9672, + "evolution_tflops": 1.2931, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_even_tail_b2_n1792_k256_d128", + "params": { + "B": 2, + "D": 128, + "K": 256, + "N": 1792, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.9451, + "evolution_flashlib_ms": 0.248512, + "evolution_kernel_ms": 0.164208, + "evolution_speedup": 1.5134, + "evolution_tflops": 1.4304, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "tail_even_b2_n2048_k2048_d128", + "params": { + "B": 2, + "D": 128, + "K": 2048, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 3.8606, + "evolution_flashlib_ms": 0.556255, + "evolution_kernel_ms": 0.186303, + "evolution_speedup": 2.9858, + "evolution_tflops": 11.5268, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "random_legal_b2_n4736_k512_d128", + "params": { + "B": 2, + "D": 128, + "K": 512, + "N": 4736, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 3.587, + "evolution_flashlib_ms": 0.346112, + "evolution_kernel_ms": 0.146688, + "evolution_speedup": 2.3595, + "evolution_tflops": 8.4636, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_large_b2_n65664_k512_d128", + "params": { + "B": 2, + "D": 128, + "K": 512, + "N": 65664, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 40.2273, + "evolution_flashlib_ms": 0.427904, + "evolution_kernel_ms": 0.17792, + "evolution_speedup": 2.405, + "evolution_tflops": 96.7478, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "small_low_b3_n256_k512_d128", + "params": { + "B": 3, + "D": 128, + "K": 512, + "N": 256, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1812, + "evolution_flashlib_ms": 0.555391, + "evolution_kernel_ms": 0.145791, + "evolution_speedup": 3.8095, + "evolution_tflops": 0.6905, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_first_even_b3_n1536_k512_d128", + "params": { + "B": 3, + "D": 128, + "K": 512, + "N": 1536, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.7141, + "evolution_flashlib_ms": 0.352351, + "evolution_kernel_ms": 0.166016, + "evolution_speedup": 2.1224, + "evolution_tflops": 3.6381, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_mid_odd_b3_n2432_k512_d128", + "params": { + "B": 3, + "D": 128, + "K": 512, + "N": 2432, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 3.0254, + "evolution_flashlib_ms": 0.316095, + "evolution_kernel_ms": 0.144, + "evolution_speedup": 2.1951, + "evolution_tflops": 6.641, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "random_legal_b3_n2560_k1536_d128", + "params": { + "B": 3, + "D": 128, + "K": 1536, + "N": 2560, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 7.0022, + "evolution_flashlib_ms": 0.43128, + "evolution_kernel_ms": 0.179008, + "evolution_speedup": 2.4093, + "evolution_tflops": 16.8702, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "random_b3_n6144_k1280_d128", + "params": { + "B": 3, + "D": 128, + "K": 1280, + "N": 6144, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 14.1328, + "evolution_flashlib_ms": 0.427359, + "evolution_kernel_ms": 0.17632, + "evolution_speedup": 2.4238, + "evolution_tflops": 34.2548, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "small_b4_n128_k512_d128", + "params": { + "B": 4, + "D": 128, + "K": 512, + "N": 128, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1386, + "evolution_flashlib_ms": 0.484159, + "evolution_kernel_ms": 0.144912, + "evolution_speedup": 3.3411, + "evolution_tflops": 0.4631, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "random_legal_b4_n640_k1024_d128", + "params": { + "B": 4, + "D": 128, + "K": 1024, + "N": 640, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.2317, + "evolution_flashlib_ms": 0.544863, + "evolution_kernel_ms": 0.148095, + "evolution_speedup": 3.6791, + "evolution_tflops": 4.5315, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "boundary_b4_n1024_k512_d128", + "params": { + "B": 4, + "D": 128, + "K": 512, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.3605, + "evolution_flashlib_ms": 0.394623, + "evolution_kernel_ms": 0.144512, + "evolution_speedup": 2.7307, + "evolution_tflops": 3.7151, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_odd_b4_n1664_k256_d128", + "params": { + "B": 4, + "D": 128, + "K": 256, + "N": 1664, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.4791, + "evolution_flashlib_ms": 0.294912, + "evolution_kernel_ms": 0.145728, + "evolution_speedup": 2.0237, + "evolution_tflops": 2.9933, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_widek_b4_n3072_k3072_d128", + "params": { + "B": 4, + "D": 128, + "K": 3072, + "N": 3072, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 15.4254, + "evolution_flashlib_ms": 0.62648, + "evolution_kernel_ms": 0.19712, + "evolution_speedup": 3.1782, + "evolution_tflops": 49.0242, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "random_legal_b5_n512_k512_d128", + "params": { + "B": 5, + "D": 128, + "K": 512, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.7414, + "evolution_flashlib_ms": 0.452576, + "evolution_kernel_ms": 0.14368, + "evolution_speedup": 3.1499, + "evolution_tflops": 2.3354, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "random_b5_n3456_k512_d128", + "params": { + "B": 5, + "D": 128, + "K": 512, + "N": 3456, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 7.5675, + "evolution_flashlib_ms": 0.299296, + "evolution_kernel_ms": 0.145872, + "evolution_speedup": 2.0518, + "evolution_tflops": 15.5268, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "small_cap_b6_n1024_k256_d128", + "params": { + "B": 6, + "D": 128, + "K": 256, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.141, + "evolution_flashlib_ms": 0.352896, + "evolution_kernel_ms": 0.144144, + "evolution_speedup": 2.4482, + "evolution_tflops": 2.7934, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_large_b6_n65536_k512_d128", + "params": { + "B": 6, + "D": 128, + "K": 512, + "N": 65536, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 57.2499, + "evolution_flashlib_ms": 0.900256, + "evolution_kernel_ms": 0.253025, + "evolution_speedup": 3.558, + "evolution_tflops": 203.6937, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "small_boundary_b7_n896_k256_d128", + "params": { + "B": 7, + "D": 128, + "K": 256, + "N": 896, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.2543, + "evolution_flashlib_ms": 0.327696, + "evolution_kernel_ms": 0.14368, + "evolution_speedup": 2.2807, + "evolution_tflops": 2.8608, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "random_b7_n1024_k4096_d128", + "params": { + "B": 7, + "D": 128, + "K": 4096, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 8.1063, + "evolution_flashlib_ms": 0.9272, + "evolution_kernel_ms": 0.20928, + "evolution_speedup": 4.4304, + "evolution_tflops": 35.9145, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_b8_n3712_k2048_d128", + "params": { + "B": 8, + "D": 128, + "K": 2048, + "N": 3712, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 25.0897, + "evolution_flashlib_ms": 0.620543, + "evolution_kernel_ms": 0.171536, + "evolution_speedup": 3.6176, + "evolution_tflops": 90.764, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_b8_n4096_k256_d128", + "params": { + "B": 8, + "D": 128, + "K": 256, + "N": 4096, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 8.9324, + "evolution_flashlib_ms": 0.240416, + "evolution_kernel_ms": 0.164624, + "evolution_speedup": 1.4604, + "evolution_tflops": 13.0448, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "mid_b8_n8192_k512_d128", + "params": { + "B": 8, + "D": 128, + "K": 512, + "N": 8192, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 21.7357, + "evolution_flashlib_ms": 0.395199, + "evolution_kernel_ms": 0.176959, + "evolution_speedup": 2.2333, + "evolution_tflops": 48.5418, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "random_legal_b9_n3840_k1792_d128", + "params": { + "B": 9, + "D": 128, + "K": 1792, + "N": 3840, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 26.2993, + "evolution_flashlib_ms": 0.602848, + "evolution_kernel_ms": 0.186464, + "evolution_speedup": 3.2331, + "evolution_tflops": 85.027, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "small_b12_n640_k256_d128", + "params": { + "B": 12, + "D": 128, + "K": 256, + "N": 640, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.3627, + "evolution_flashlib_ms": 0.36936, + "evolution_kernel_ms": 0.144128, + "evolution_speedup": 2.5627, + "evolution_tflops": 3.4921, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "small_b16_n1024_k512_d128", + "params": { + "B": 16, + "D": 128, + "K": 512, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 5.3198, + "evolution_flashlib_ms": 0.40368, + "evolution_kernel_ms": 0.1464, + "evolution_speedup": 2.7574, + "evolution_tflops": 14.6686, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "large_b16_n32768_k1024_d128", + "params": { + "B": 16, + "D": 128, + "K": 1024, + "N": 32768, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 65.9244, + "evolution_flashlib_ms": 2.084797, + "evolution_kernel_ms": 0.363695, + "evolution_speedup": 5.7323, + "evolution_tflops": 377.8962, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paper_b32_n75776_k1024_d128", + "params": { + "B": 32, + "D": 128, + "K": 1024, + "N": 75776, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 75.3887, + "evolution_flashlib_ms": 8.431702, + "evolution_kernel_ms": 0.951326, + "evolution_speedup": 8.8631, + "evolution_tflops": 668.1781, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d160_paired_b1_n2048_k2048_d160", + "params": { + "B": 1, + "D": 160, + "K": 2048, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.3001, + "evolution_flashlib_ms": 0.583535, + "evolution_kernel_ms": 0.188448, + "evolution_speedup": 3.0965, + "evolution_tflops": 7.1223, + "route": "d160_padded_single_repeated_mma_v2", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d160_fallback_b2_n2432_k1024_d160", + "params": { + "B": 2, + "D": 160, + "K": 1024, + "N": 2432, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 3.798, + "evolution_flashlib_ms": 0.419648, + "evolution_kernel_ms": 0.18832, + "evolution_speedup": 2.2284, + "evolution_tflops": 8.4634, + "route": "d160_padded_single_repeated_mma_v2", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d192_fallback_b1_n2176_k1024_d192", + "params": { + "B": 1, + "D": 192, + "K": 1024, + "N": 2176, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.0726, + "evolution_flashlib_ms": 0.412832, + "evolution_kernel_ms": 0.149952, + "evolution_speedup": 2.7531, + "evolution_tflops": 5.7061, + "route": "d192_single_repeated_mma_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d192_small_b2_n1024_k512_d192", + "params": { + "B": 2, + "D": 192, + "K": 512, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.8948, + "evolution_flashlib_ms": 0.449999, + "evolution_kernel_ms": 0.146657, + "evolution_speedup": 3.0684, + "evolution_tflops": 2.7455, + "route": "d192_single_repeated_mma_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d192_paired_b2_n2048_k2048_d192", + "params": { + "B": 2, + "D": 192, + "K": 2048, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 4.6837, + "evolution_flashlib_ms": 0.687759, + "evolution_kernel_ms": 0.310336, + "evolution_speedup": 2.2162, + "evolution_tflops": 10.3798, + "route": "d192_paired_repeated_mma_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d192_large_b4_n32768_k1024_d192", + "params": { + "B": 4, + "D": 192, + "K": 1024, + "N": 32768, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 56.4296, + "evolution_flashlib_ms": 0.913344, + "evolution_kernel_ms": 0.221536, + "evolution_speedup": 4.1228, + "evolution_tflops": 232.6466, + "route": "d192_paired_repeated_mma_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d256_hugek_b1_n512_k8192_d256", + "params": { + "B": 1, + "D": 256, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0389, + "evolution_flashlib_ms": 55.197531, + "evolution_kernel_ms": 0.222944, + "evolution_speedup": 247.5847, + "evolution_tflops": 9.6324, + "route": "d256_single_repeated_mma_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d256_small_b1_n1024_k512_d256", + "params": { + "B": 1, + "D": 256, + "K": 512, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.6249, + "evolution_flashlib_ms": 0.429567, + "evolution_kernel_ms": 0.148224, + "evolution_speedup": 2.8981, + "evolution_tflops": 1.811, + "route": "d256_single_repeated_mma_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d256_paired_b1_n4096_k4096_d256", + "params": { + "B": 1, + "D": 256, + "K": 4096, + "N": 4096, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1557, + "evolution_flashlib_ms": 55.152968, + "evolution_kernel_ms": 0.179455, + "evolution_speedup": 307.3359, + "evolution_tflops": 47.8668, + "route": "d256_single_repeated_mma_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d256_fallback_b2_n2432_k2048_d256", + "params": { + "B": 2, + "D": 256, + "K": 2048, + "N": 2432, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 8.2786, + "evolution_flashlib_ms": 0.616079, + "evolution_kernel_ms": 0.16352, + "evolution_speedup": 3.7676, + "evolution_tflops": 31.1905, + "route": "d256_single_repeated_mma_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d320_paired_b1_n2048_k4096_d320", + "params": { + "B": 1, + "D": 320, + "K": 4096, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1154, + "evolution_flashlib_ms": 46.53602, + "evolution_kernel_ms": 0.23344, + "evolution_speedup": 199.349, + "evolution_tflops": 22.9982, + "route": "highd_splitd_single_tile_6fcf_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d320_large_b2_n16384_k1024_d320", + "params": { + "B": 2, + "D": 320, + "K": 1024, + "N": 16384, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 36.5995, + "evolution_flashlib_ms": 0.586752, + "evolution_kernel_ms": 0.19096, + "evolution_speedup": 3.0726, + "evolution_tflops": 112.4573, + "route": "highd_splitd_single_tile_6fcf_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d384_hugek_b1_n768_k8192_d384", + "params": { + "B": 1, + "D": 384, + "K": 8192, + "N": 768, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1211, + "evolution_flashlib_ms": 39.895303, + "evolution_kernel_ms": 0.271664, + "evolution_speedup": 146.8553, + "evolution_tflops": 17.7861, + "route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d384_small_b1_n896_k512_d384", + "params": { + "B": 1, + "D": 384, + "K": 512, + "N": 896, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.6757, + "evolution_flashlib_ms": 0.521424, + "evolution_kernel_ms": 0.15824, + "evolution_speedup": 3.2951, + "evolution_tflops": 2.2265, + "route": "highd_splitd_single_tile_6fcf_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d384_paired_b1_n2048_k4096_d384", + "params": { + "B": 1, + "D": 384, + "K": 4096, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1674, + "evolution_flashlib_ms": 38.477384, + "evolution_kernel_ms": 0.252672, + "evolution_speedup": 152.2819, + "evolution_tflops": 25.4973, + "route": "highd_splitd_single_tile_6fcf_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d384_fallback_b3_n3456_k1024_d384", + "params": { + "B": 3, + "D": 384, + "K": 1024, + "N": 3456, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 13.6962, + "evolution_flashlib_ms": 0.595328, + "evolution_kernel_ms": 0.174575, + "evolution_speedup": 3.4101, + "evolution_tflops": 46.706, + "route": "highd_splitd_single_tile_6fcf_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d448_hugek_b1_n512_k8192_d448", + "params": { + "B": 1, + "D": 448, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0923, + "evolution_flashlib_ms": 40.698431, + "evolution_kernel_ms": 0.262959, + "evolution_speedup": 154.771, + "evolution_tflops": 14.2916, + "route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d448_paired_b1_n2048_k4096_d448", + "params": { + "B": 1, + "D": 448, + "K": 4096, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1943, + "evolution_flashlib_ms": 38.67485, + "evolution_kernel_ms": 0.269632, + "evolution_speedup": 143.4357, + "evolution_tflops": 27.8757, + "route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d512_hugek_b1_n512_k8192_d512", + "params": { + "B": 1, + "D": 512, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1326, + "evolution_flashlib_ms": 32.378557, + "evolution_kernel_ms": 0.266288, + "evolution_speedup": 121.5925, + "evolution_tflops": 16.1291, + "route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d512_small_b1_n1024_k512_d512", + "params": { + "B": 1, + "D": 512, + "K": 512, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.969, + "evolution_flashlib_ms": 0.554048, + "evolution_kernel_ms": 0.162688, + "evolution_speedup": 3.4056, + "evolution_tflops": 3.3, + "route": "highd_splitd_single_tile_6fcf_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d512_paired_b1_n2048_k4096_d512", + "params": { + "B": 1, + "D": 512, + "K": 4096, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.281, + "evolution_flashlib_ms": 30.573365, + "evolution_kernel_ms": 0.270976, + "evolution_speedup": 112.8268, + "evolution_tflops": 31.7, + "route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d512_fallback_b2_n2432_k2048_d512", + "params": { + "B": 2, + "D": 512, + "K": 2048, + "N": 2432, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 11.9237, + "evolution_flashlib_ms": 0.855488, + "evolution_kernel_ms": 0.220319, + "evolution_speedup": 3.883, + "evolution_tflops": 46.299, + "route": "highd_splitd_single_tile_6fcf_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "post_d895_d16_b4_n32768_k1024_d16", + "params": { + "B": 4, + "D": 16, + "K": 1024, + "N": 32768, + "dtype": "bfloat16", + "seed": 21602, + "source": "near_floor_microdim" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 12.993, + "evolution_flashlib_ms": 0.330559, + "evolution_kernel_ms": 0.216591, + "evolution_route": "microdim_hybrid_9c0d_v1", + "evolution_speedup": 1.5262, + "evolution_tflops": 19.8298 + } + }, + { + "label": "post_d895_d16_b8_n65536_k512_d16", + "params": { + "B": 8, + "D": 16, + "K": 512, + "N": 65536, + "dtype": "bfloat16", + "seed": 21601, + "source": "near_floor_microdim" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 19.3119, + "evolution_flashlib_ms": 0.444799, + "evolution_kernel_ms": 0.295359, + "evolution_route": "microdim_hybrid_9c0d_v1", + "evolution_speedup": 1.506, + "evolution_tflops": 29.083 + } + }, + { + "label": "post_d895_d32_b4_n32768_k1024_d32", + "params": { + "B": 4, + "D": 32, + "K": 1024, + "N": 32768, + "dtype": "bfloat16", + "seed": 23202, + "source": "near_floor_microdim" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 23.8017, + "evolution_flashlib_ms": 0.360896, + "evolution_kernel_ms": 0.218816, + "evolution_route": "microdim_hybrid_9c0d_v1", + "evolution_speedup": 1.6493, + "evolution_tflops": 39.2564 + } + }, + { + "label": "post_d895_d32_b8_n65536_k512_d32", + "params": { + "B": 8, + "D": 32, + "K": 512, + "N": 65536, + "dtype": "bfloat16", + "seed": 23201, + "source": "near_floor_microdim" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 34.3587, + "evolution_flashlib_ms": 0.500015, + "evolution_kernel_ms": 0.301439, + "evolution_route": "microdim_hybrid_9c0d_v1", + "evolution_speedup": 1.6588, + "evolution_tflops": 56.9928 + } + }, + { + "label": "adjacent_3328_d48_small_boundary_b1_n256_k256_d48", + "params": { + "B": 1, + "D": 48, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "seed": 3328481, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0145, + "evolution_flashlib_ms": 0.434063, + "evolution_kernel_ms": 0.168656, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.5737, + "evolution_tflops": 0.0373 + } + }, + { + "label": "post_d895_d48_b1_n512_k8192_d48", + "params": { + "B": 1, + "D": 48, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 24803, + "source": "high_k_low_n" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.275, + "evolution_flashlib_ms": 1.464301, + "evolution_kernel_ms": 0.21952, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 6.6705, + "evolution_tflops": 1.8342 + } + }, + { + "label": "adjacent_d9d5_d48_highk_heldout_b2_n640_k4096_d48", + "params": { + "B": 2, + "D": 48, + "K": 4096, + "N": 640, + "dtype": "bfloat16", + "seed": 9504801, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.5494, + "evolution_flashlib_ms": 0.916111, + "evolution_kernel_ms": 0.19392, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 4.7242, + "evolution_tflops": 2.5955 + } + }, + { + "label": "adjacent_a2f8_d48_k1024_boundary_b2_n768_k1024_d48", + "params": { + "B": 2, + "D": 48, + "K": 1024, + "N": 768, + "dtype": "bfloat16", + "seed": 1028481, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.3038, + "evolution_flashlib_ms": 0.497055, + "evolution_kernel_ms": 0.173184, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.8701, + "evolution_tflops": 0.8719 + } + }, + { + "label": "adjacent_8f09_d48_medium_tail_b2_n1792_k512_d48", + "params": { + "B": 2, + "D": 48, + "K": 512, + "N": 1792, + "dtype": "bfloat16", + "seed": 8090481, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.556, + "evolution_flashlib_ms": 0.316848, + "evolution_kernel_ms": 0.169375, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.8707, + "evolution_tflops": 1.0401 + } + }, + { + "label": "post_d895_d48_b2_n2048_k512_d48", + "params": { + "B": 2, + "D": 48, + "K": 512, + "N": 2048, + "dtype": "bfloat16", + "seed": 24801, + "source": "new_lowmid_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.6574, + "evolution_flashlib_ms": 0.30624, + "evolution_kernel_ms": 0.166656, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.8376, + "evolution_tflops": 1.208 + } + }, + { + "label": "adjacent_5600_d48_low_n_boundary_b3_n1536_k256_d48", + "params": { + "B": 3, + "D": 48, + "K": 256, + "N": 1536, + "dtype": "bfloat16", + "seed": 5600481, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.4029, + "evolution_flashlib_ms": 0.281087, + "evolution_kernel_ms": 0.169183, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.6614, + "evolution_tflops": 0.6694 + } + }, + { + "label": "adjacent_c44f_d48_midk_boundary_b3_n2688_k768_d48", + "params": { + "B": 3, + "D": 48, + "K": 768, + "N": 2688, + "dtype": "bfloat16", + "seed": 4404801, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 1.9182, + "evolution_flashlib_ms": 0.309952, + "evolution_kernel_ms": 0.174383, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.7774, + "evolution_tflops": 3.4094 + } + }, + { + "label": "adjacent_ef00_d48_midn_boundary_b3_n3456_k512_d48", + "params": { + "B": 3, + "D": 48, + "K": 512, + "N": 3456, + "dtype": "bfloat16", + "seed": 9000481, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 1.9767, + "evolution_flashlib_ms": 0.257807, + "evolution_kernel_ms": 0.172495, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.4946, + "evolution_tflops": 2.9543 + } + }, + { + "label": "adjacent_68cf_d48_boundary_b4_n2304_k512_d48", + "params": { + "B": 4, + "D": 48, + "K": 512, + "N": 2304, + "dtype": "bfloat16", + "seed": 6804801, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 1.5408, + "evolution_flashlib_ms": 0.293984, + "evolution_kernel_ms": 0.170608, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.7232, + "evolution_tflops": 2.6551 + } + }, + { + "label": "adjacent_1d49_d48_lowk_boundary_b4_n3968_k256_d48", + "params": { + "B": 4, + "D": 48, + "K": 256, + "N": 3968, + "dtype": "bfloat16", + "seed": 1490481, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 1.7051, + "evolution_flashlib_ms": 0.228768, + "evolution_kernel_ms": 0.17192, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.3307, + "evolution_tflops": 2.2689 + } + }, + { + "label": "post_d895_d48_b4_n8192_k1024_d48", + "params": { + "B": 4, + "D": 48, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "seed": 24802, + "source": "new_lowmid_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 9.5879, + "evolution_flashlib_ms": 0.335967, + "evolution_kernel_ms": 0.179807, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.8685, + "evolution_tflops": 17.9149 + } + }, + { + "label": "adjacent_9ca0_d48_guard_boundary_b6_n12288_k512_d48", + "params": { + "B": 6, + "D": 48, + "K": 512, + "N": 12288, + "dtype": "bfloat16", + "seed": 924805, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 11.385, + "evolution_flashlib_ms": 0.318304, + "evolution_kernel_ms": 0.182336, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.7457, + "evolution_tflops": 19.8747 + } + }, + { + "label": "post_d895_d112_b1_n256_k256_d112", + "params": { + "B": 1, + "D": 112, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "seed": 211204, + "source": "lowmid_small_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0353, + "evolution_flashlib_ms": 0.416031, + "evolution_kernel_ms": 0.166575, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.4976, + "evolution_tflops": 0.0881 + } + }, + { + "label": "adjacent_8f09_d112_small_highk_b1_n384_k4096_d112", + "params": { + "B": 1, + "D": 112, + "K": 4096, + "N": 384, + "dtype": "bfloat16", + "seed": 8091121, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.2782, + "evolution_flashlib_ms": 1.266206, + "evolution_kernel_ms": 0.195296, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 6.4835, + "evolution_tflops": 1.804 + } + }, + { + "label": "post_d895_d112_b1_n512_k8192_d112", + "params": { + "B": 1, + "D": 112, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 211203, + "source": "high_k_low_n" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.5425, + "evolution_flashlib_ms": 1.731693, + "evolution_kernel_ms": 0.294304, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 5.884, + "evolution_tflops": 3.1924 + } + }, + { + "label": "adjacent_1d49_d112_highk_tail_b1_n1408_k4096_d112", + "params": { + "B": 1, + "D": 112, + "K": 4096, + "N": 1408, + "dtype": "bfloat16", + "seed": 1491121, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 1.4908, + "evolution_flashlib_ms": 0.866559, + "evolution_kernel_ms": 0.196192, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 4.4169, + "evolution_tflops": 6.5846 + } + }, + { + "label": "adjacent_3328_d112_highk_low_n_b2_n512_k8192_d112", + "params": { + "B": 2, + "D": 112, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 3328112, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 1.0686, + "evolution_flashlib_ms": 1.758429, + "evolution_kernel_ms": 0.295872, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 5.9432, + "evolution_tflops": 6.3509 + } + }, + { + "label": "post_d895_d112_b2_n2048_k512_d112", + "params": { + "B": 2, + "D": 112, + "K": 512, + "N": 2048, + "dtype": "bfloat16", + "seed": 211201, + "source": "new_lowmid_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 1.5026, + "evolution_flashlib_ms": 0.312623, + "evolution_kernel_ms": 0.188, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.6629, + "evolution_tflops": 2.4987 + } + }, + { + "label": "adjacent_c44f_d112_odd_tail_b2_n3200_k768_d112", + "params": { + "B": 2, + "D": 112, + "K": 768, + "N": 3200, + "dtype": "bfloat16", + "seed": 4411201, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 3.4225, + "evolution_flashlib_ms": 0.321696, + "evolution_kernel_ms": 0.173088, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.8586, + "evolution_tflops": 6.361 + } + }, + { + "label": "adjacent_d9d5_d112_highk_request_b3_n768_k8192_d112", + "params": { + "B": 3, + "D": 112, + "K": 8192, + "N": 768, + "dtype": "bfloat16", + "seed": 9511201, + "source": "request_specific" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.7846, + "evolution_flashlib_ms": 1.518318, + "evolution_kernel_ms": 0.287872, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 5.2743, + "evolution_tflops": 14.6866 + } + }, + { + "label": "adjacent_9ca0_d112_tail_div_b3_n3840_k768_d112", + "params": { + "B": 3, + "D": 112, + "K": 768, + "N": 3840, + "dtype": "bfloat16", + "seed": 921125, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 5.9472, + "evolution_flashlib_ms": 0.333232, + "evolution_kernel_ms": 0.194592, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.7125, + "evolution_tflops": 10.1844 + } + }, + { + "label": "adjacent_a2f8_d112_lowk_tail_b4_n3456_k256_d112", + "params": { + "B": 4, + "D": 112, + "K": 256, + "N": 3456, + "dtype": "bfloat16", + "seed": 1028112, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 3.4264, + "evolution_flashlib_ms": 0.231359, + "evolution_kernel_ms": 0.168912, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.3697, + "evolution_tflops": 4.6931 + } + }, + { + "label": "adjacent_ef00_d112_odd_tail_b4_n3712_k1024_d112", + "params": { + "B": 4, + "D": 112, + "K": 1024, + "N": 3712, + "dtype": "bfloat16", + "seed": 9001121, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 8.805, + "evolution_flashlib_ms": 0.386799, + "evolution_kernel_ms": 0.176608, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.1902, + "evolution_tflops": 19.2844 + } + }, + { + "label": "post_d895_d112_b4_n8192_k1024_d112", + "params": { + "B": 4, + "D": 112, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "seed": 211202, + "source": "new_lowmid_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 18.1068, + "evolution_flashlib_ms": 0.415103, + "evolution_kernel_ms": 0.1956, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.1222, + "evolution_tflops": 38.4264 + } + }, + { + "label": "adjacent_5600_d112_odd_tile_tail_b5_n2176_k512_d112", + "params": { + "B": 5, + "D": 112, + "K": 512, + "N": 2176, + "dtype": "bfloat16", + "seed": 5601121, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 3.6867, + "evolution_flashlib_ms": 0.338463, + "evolution_kernel_ms": 0.171519, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.9733, + "evolution_tflops": 7.275 + } + }, + { + "label": "adjacent_68cf_d112_tail_b5_n2944_k512_d112", + "params": { + "B": 5, + "D": 112, + "K": 512, + "N": 2944, + "dtype": "bfloat16", + "seed": 6811201, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 5.7016, + "evolution_flashlib_ms": 0.296095, + "evolution_kernel_ms": 0.173023, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.7113, + "evolution_tflops": 9.7571 + } + }, + { + "label": "post_d895_d128_paired_b2_n262144_k256_d128", + "params": { + "B": 2, + "D": 128, + "K": 256, + "N": 262144, + "dtype": "bfloat16", + "seed": 212802, + "source": "near_floor_paired_d128" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 51.8017, + "evolution_flashlib_ms": 0.663294, + "evolution_kernel_ms": 0.24464, + "evolution_route": "paired_large_v15", + "evolution_speedup": 2.7113, + "evolution_tflops": 140.4505 + } + }, + { + "label": "post_d895_d128_fallback_b3_n1920_k256_d128", + "params": { + "B": 3, + "D": 128, + "K": 256, + "N": 1920, + "dtype": "bfloat16", + "seed": 212803, + "source": "near_floor_fallback_d128" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 1.3402, + "evolution_flashlib_ms": 0.281664, + "evolution_kernel_ms": 0.1436, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 1.9614, + "evolution_tflops": 2.6287 + } + }, + { + "label": "adjacent_8f09_d128_fallback_neighbor_b4_n4480_k512_d128", + "params": { + "B": 4, + "D": 128, + "K": 512, + "N": 4480, + "dtype": "bfloat16", + "seed": 8091281, + "source": "guard_miss_fallback" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 7.2152, + "evolution_flashlib_ms": 0.325535, + "evolution_kernel_ms": 0.145856, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 2.2319, + "evolution_tflops": 16.1036 + } + }, + { + "label": "adjacent_ef00_d128_forced_fallback_b4_n7552_k512_d128", + "params": { + "B": 4, + "D": 128, + "K": 512, + "N": 7552, + "dtype": "bfloat16", + "seed": 9001281, + "source": "forced_fallback" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 12.6084, + "evolution_flashlib_ms": 0.314032, + "evolution_kernel_ms": 0.153023, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 2.0522, + "evolution_tflops": 25.8747 + } + }, + { + "label": "post_d895_d128_fallback_b5_n2176_k512_d128", + "params": { + "B": 5, + "D": 128, + "K": 512, + "N": 2176, + "dtype": "bfloat16", + "seed": 212804, + "source": "near_floor_fallback_d128" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 4.6116, + "evolution_flashlib_ms": 0.309232, + "evolution_kernel_ms": 0.14432, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 2.1427, + "evolution_tflops": 9.8813 + } + }, + { + "label": "adjacent_d9d5_d128_guard_miss_fallback_b5_n6016_k1024_d128", + "params": { + "B": 5, + "D": 128, + "K": 1024, + "N": 6016, + "dtype": "bfloat16", + "seed": 9512801, + "source": "guard_miss_fallback" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 17.1342, + "evolution_flashlib_ms": 0.460207, + "evolution_kernel_ms": 0.159104, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 2.8925, + "evolution_tflops": 49.5606 + } + }, + { + "label": "adjacent_1d49_d128_forced_fallback_b5_n7296_k512_d128", + "params": { + "B": 5, + "D": 128, + "K": 512, + "N": 7296, + "dtype": "bfloat16", + "seed": 1491281, + "source": "forced_fallback" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 14.0693, + "evolution_flashlib_ms": 0.339855, + "evolution_kernel_ms": 0.152639, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 2.2265, + "evolution_tflops": 31.3255 + } + }, + { + "label": "adjacent_c44f_d128_forced_fallback_b6_n6272_k512_d128", + "params": { + "B": 6, + "D": 128, + "K": 512, + "N": 6272, + "dtype": "bfloat16", + "seed": 4412801, + "source": "forced_fallback" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 14.7956, + "evolution_flashlib_ms": 0.333376, + "evolution_kernel_ms": 0.152256, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 2.1896, + "evolution_tflops": 32.3961 + } + }, + { + "label": "adjacent_a2f8_d128_odd_fallback_b6_n8576_k512_d128", + "params": { + "B": 6, + "D": 128, + "K": 512, + "N": 8576, + "dtype": "bfloat16", + "seed": 1028128, + "source": "guard_miss_fallback" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 19.2761, + "evolution_flashlib_ms": 0.349887, + "evolution_kernel_ms": 0.158688, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 2.2049, + "evolution_tflops": 42.5013 + } + }, + { + "label": "post_d895_d128_fallback_b7_n2432_k1024_d128", + "params": { + "B": 7, + "D": 128, + "K": 1024, + "N": 2432, + "dtype": "bfloat16", + "seed": 212805, + "source": "near_floor_fallback_d128" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 11.0028, + "evolution_flashlib_ms": 0.4056, + "evolution_kernel_ms": 0.147936, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 2.7417, + "evolution_tflops": 30.1667 + } + }, + { + "label": "adjacent_68cf_d128_forced_fallback_b7_n6016_k512_d128", + "params": { + "B": 7, + "D": 128, + "K": 512, + "N": 6016, + "dtype": "bfloat16", + "seed": 6812801, + "source": "forced_fallback" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 14.2803, + "evolution_flashlib_ms": 0.386527, + "evolution_kernel_ms": 0.156832, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 2.4646, + "evolution_tflops": 35.195 + } + }, + { + "label": "adjacent_9ca0_d128_exact_guard_neighbor_b8_n8064_k256_d128", + "params": { + "B": 8, + "D": 128, + "K": 256, + "N": 8064, + "dtype": "bfloat16", + "seed": 912806, + "source": "guard_miss_fallback" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 14.5156, + "evolution_flashlib_ms": 0.291264, + "evolution_kernel_ms": 0.155168, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 1.8771, + "evolution_tflops": 27.247 + } + }, + { + "label": "post_d895_d128_paired_b8_n8192_k256_d128", + "params": { + "B": 8, + "D": 128, + "K": 256, + "N": 8192, + "dtype": "bfloat16", + "seed": 212801, + "source": "near_floor_paired_d128" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 14.6751, + "evolution_flashlib_ms": 0.292671, + "evolution_kernel_ms": 0.15424, + "evolution_route": "d128_even_near_floor_v10_repair", + "evolution_speedup": 1.8975, + "evolution_tflops": 27.846 + } + }, + { + "label": "adjacent_3328_d128_exact_guard_k_neighbor_b8_n8192_k512_d128", + "params": { + "B": 8, + "D": 128, + "K": 512, + "N": 8192, + "dtype": "bfloat16", + "seed": 3328128, + "source": "guard_miss_fallback" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 21.7692, + "evolution_flashlib_ms": 0.394592, + "evolution_kernel_ms": 0.177472, + "evolution_route": "paired_large_v15", + "evolution_speedup": 2.2234, + "evolution_tflops": 48.4016 + } + }, + { + "label": "adjacent_5600_d128_fallback_neighbor_b8_n8320_k256_d128", + "params": { + "B": 8, + "D": 128, + "K": 256, + "N": 8320, + "dtype": "bfloat16", + "seed": 5601281, + "source": "guard_miss_fallback" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 14.6812, + "evolution_flashlib_ms": 0.297119, + "evolution_kernel_ms": 0.155872, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 1.9062, + "evolution_tflops": 27.985 + } + }, + { + "label": "post_d895_d144_b1_n256_k256_d144", + "params": { + "B": 1, + "D": 144, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "seed": 214401, + "source": "tailpad_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0351, + "evolution_flashlib_ms": 0.537279, + "evolution_kernel_ms": 0.167728, + "evolution_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_speedup": 3.2033, + "evolution_tflops": 0.1125 + } + }, + { + "label": "post_d895_d144_b1_n512_k8192_d144", + "params": { + "B": 1, + "D": 144, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 214404, + "source": "high_k_low_n" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.6035, + "evolution_flashlib_ms": 2.001693, + "evolution_kernel_ms": 0.302528, + "evolution_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_speedup": 6.6166, + "evolution_tflops": 3.9929 + } + }, + { + "label": "post_d895_d144_b2_n2048_k1024_d144", + "params": { + "B": 2, + "D": 144, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "seed": 214402, + "source": "tailpad_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.5941, + "evolution_flashlib_ms": 0.465663, + "evolution_kernel_ms": 0.192896, + "evolution_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_speedup": 2.4141, + "evolution_tflops": 6.2622 + } + }, + { + "label": "post_d895_d144_b4_n8192_k1024_d144", + "params": { + "B": 4, + "D": 144, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "seed": 214403, + "source": "tailpad_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 20.3965, + "evolution_flashlib_ms": 0.473791, + "evolution_kernel_ms": 0.195104, + "evolution_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_speedup": 2.4284, + "evolution_tflops": 49.5309 + } + }, + { + "label": "post_d895_d176_b1_n256_k256_d176", + "params": { + "B": 1, + "D": 176, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "seed": 217601, + "source": "tailpad_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0453, + "evolution_flashlib_ms": 0.509327, + "evolution_kernel_ms": 0.169792, + "evolution_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_speedup": 2.9997, + "evolution_tflops": 0.1359 + } + }, + { + "label": "post_d895_d176_b1_n512_k8192_d176", + "params": { + "B": 1, + "D": 176, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 217604, + "source": "high_k_low_n" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.72, + "evolution_flashlib_ms": 2.050605, + "evolution_kernel_ms": 0.301759, + "evolution_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_speedup": 6.7955, + "evolution_tflops": 4.8926 + } + }, + { + "label": "post_d895_d176_b2_n2048_k1024_d176", + "params": { + "B": 2, + "D": 176, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "seed": 217602, + "source": "tailpad_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 3.0356, + "evolution_flashlib_ms": 0.486367, + "evolution_kernel_ms": 0.191839, + "evolution_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_speedup": 2.5353, + "evolution_tflops": 7.696 + } + }, + { + "label": "post_d895_d176_b4_n8192_k1024_d176", + "params": { + "B": 4, + "D": 176, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "seed": 217603, + "source": "tailpad_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 23.9979, + "evolution_flashlib_ms": 0.492175, + "evolution_kernel_ms": 0.195712, + "evolution_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_speedup": 2.5148, + "evolution_tflops": 60.3497 + } + }, + { + "label": "post_d895_d224_b1_n256_k256_d224", + "params": { + "B": 1, + "D": 224, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "seed": 222404, + "source": "medium_small_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0611, + "evolution_flashlib_ms": 0.480639, + "evolution_kernel_ms": 0.170496, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.8191, + "evolution_tflops": 0.1722 + } + }, + { + "label": "post_d895_d224_b1_n512_k8192_d224", + "params": { + "B": 1, + "D": 224, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 222403, + "source": "high_k_low_n" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.9081, + "evolution_flashlib_ms": 2.069245, + "evolution_kernel_ms": 0.236528, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 8.7484, + "evolution_tflops": 7.9443 + } + }, + { + "label": "post_d895_d224_b2_n2048_k1024_d224", + "params": { + "B": 2, + "D": 224, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "seed": 222401, + "source": "new_medium_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 3.9885, + "evolution_flashlib_ms": 0.47112, + "evolution_kernel_ms": 0.176704, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.6662, + "evolution_tflops": 10.6339 + } + }, + { + "label": "adjacent_d9d5_d224_random_midk_b2_n2944_k1280_d224", + "params": { + "B": 2, + "D": 224, + "K": 1280, + "N": 2944, + "dtype": "bfloat16", + "seed": 9522401, + "source": "random_legal" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 6.9153, + "evolution_flashlib_ms": 0.488255, + "evolution_kernel_ms": 0.180384, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.7068, + "evolution_tflops": 18.7179 + } + }, + { + "label": "adjacent_9ca0_d224_guard_overlap_b2_n4096_k256_d224", + "params": { + "B": 2, + "D": 224, + "K": 256, + "N": 4096, + "dtype": "bfloat16", + "seed": 922405, + "source": "guard_overlap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 3.6395, + "evolution_flashlib_ms": 0.258143, + "evolution_kernel_ms": 0.1712, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.5078, + "evolution_tflops": 5.4879 + } + }, + { + "label": "adjacent_ef00_d224_midn_overlap_b2_n5376_k512_d224", + "params": { + "B": 2, + "D": 224, + "K": 512, + "N": 5376, + "dtype": "bfloat16", + "seed": 9002241, + "source": "guard_overlap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 6.6916, + "evolution_flashlib_ms": 0.36856, + "evolution_kernel_ms": 0.177808, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.0728, + "evolution_tflops": 13.8703 + } + }, + { + "label": "adjacent_a2f8_d224_highn_overlap_b2_n6144_k768_d224", + "params": { + "B": 2, + "D": 224, + "K": 768, + "N": 6144, + "dtype": "bfloat16", + "seed": 1028224, + "source": "guard_overlap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 9.3487, + "evolution_flashlib_ms": 0.45224, + "evolution_kernel_ms": 0.176256, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.5658, + "evolution_tflops": 23.987 + } + }, + { + "label": "adjacent_5600_d224_k512_overlap_b3_n3072_k512_d224", + "params": { + "B": 3, + "D": 224, + "K": 512, + "N": 3072, + "dtype": "bfloat16", + "seed": 5602241, + "source": "guard_overlap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 6.148, + "evolution_flashlib_ms": 0.343839, + "evolution_kernel_ms": 0.174928, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.9656, + "evolution_tflops": 12.0846 + } + }, + { + "label": "adjacent_3328_d224_tail_div_b3_n3840_k512_d224", + "params": { + "B": 3, + "D": 224, + "K": 512, + "N": 3840, + "dtype": "bfloat16", + "seed": 3328224, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 8.3041, + "evolution_flashlib_ms": 0.318207, + "evolution_kernel_ms": 0.174895, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.8194, + "evolution_tflops": 15.1086 + } + }, + { + "label": "adjacent_68cf_d224_overlap_b3_n5120_k1024_d224", + "params": { + "B": 3, + "D": 224, + "K": 1024, + "N": 5120, + "dtype": "bfloat16", + "seed": 6822401, + "source": "guard_overlap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 14.4926, + "evolution_flashlib_ms": 0.486208, + "evolution_kernel_ms": 0.180032, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.7007, + "evolution_tflops": 39.1399 + } + }, + { + "label": "adjacent_8f09_d224_low_n_boundary_b4_n1536_k256_d224", + "params": { + "B": 4, + "D": 224, + "K": 256, + "N": 1536, + "dtype": "bfloat16", + "seed": 8092241, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.2208, + "evolution_flashlib_ms": 0.317296, + "evolution_kernel_ms": 0.173664, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.8271, + "evolution_tflops": 4.0575 + } + }, + { + "label": "adjacent_c44f_d224_overlap_b4_n4480_k512_d224", + "params": { + "B": 4, + "D": 224, + "K": 512, + "N": 4480, + "dtype": "bfloat16", + "seed": 4422401, + "source": "guard_overlap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 10.201, + "evolution_flashlib_ms": 0.402944, + "evolution_kernel_ms": 0.173008, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.329, + "evolution_tflops": 23.7585 + } + }, + { + "label": "post_d895_d224_b4_n8192_k1024_d224", + "params": { + "B": 4, + "D": 224, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "seed": 222402, + "source": "new_medium_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 27.1681, + "evolution_flashlib_ms": 0.553311, + "evolution_kernel_ms": 0.190783, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.9002, + "evolution_tflops": 78.7929 + } + }, + { + "label": "adjacent_1d49_d224_highn_overlap_b5_n5632_k768_d224", + "params": { + "B": 5, + "D": 224, + "K": 768, + "N": 5632, + "dtype": "bfloat16", + "seed": 1492241, + "source": "guard_overlap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 20.8173, + "evolution_flashlib_ms": 0.465423, + "evolution_kernel_ms": 0.186336, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.4978, + "evolution_tflops": 51.9966 + } + }, + { + "label": "post_d895_d288_b1_n256_k256_d288", + "params": { + "B": 1, + "D": 288, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "seed": 228804, + "source": "medium_small_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0649, + "evolution_flashlib_ms": 0.581567, + "evolution_kernel_ms": 0.171104, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.3989, + "evolution_tflops": 0.2206 + } + }, + { + "label": "adjacent_9ca0_d288_highk_low_n_b1_n384_k4096_d288", + "params": { + "B": 1, + "D": 288, + "K": 4096, + "N": 384, + "dtype": "bfloat16", + "seed": 928805, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0159, + "evolution_flashlib_ms": 56.807529, + "evolution_kernel_ms": 0.243872, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 232.9399, + "evolution_tflops": 3.7149 + } + }, + { + "label": "post_d895_d288_b1_n512_k8192_d288", + "params": { + "B": 1, + "D": 288, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 228803, + "source": "high_k_low_n" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0421, + "evolution_flashlib_ms": 57.444236, + "evolution_kernel_ms": 0.273615, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 209.9455, + "evolution_tflops": 8.8296 + } + }, + { + "label": "adjacent_a2f8_d288_splitk_probe_b1_n640_k4096_d288", + "params": { + "B": 1, + "D": 288, + "K": 4096, + "N": 640, + "dtype": "bfloat16", + "seed": 1028288, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0262, + "evolution_flashlib_ms": 57.6586, + "evolution_kernel_ms": 0.24416, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 236.1509, + "evolution_tflops": 6.1843 + } + }, + { + "label": "adjacent_d9d5_d288_heldout_low_n_b1_n896_k4096_d288", + "params": { + "B": 1, + "D": 288, + "K": 4096, + "N": 896, + "dtype": "bfloat16", + "seed": 9528801, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0365, + "evolution_flashlib_ms": 57.939113, + "evolution_kernel_ms": 0.2448, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 236.6794, + "evolution_tflops": 8.6353 + } + }, + { + "label": "adjacent_ef00_d288_highk_heldout_b1_n1664_k4096_d288", + "params": { + "B": 1, + "D": 288, + "K": 4096, + "N": 1664, + "dtype": "bfloat16", + "seed": 9002881, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0667, + "evolution_flashlib_ms": 58.88955, + "evolution_kernel_ms": 0.247408, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 238.0261, + "evolution_tflops": 15.868 + } + }, + { + "label": "adjacent_8f09_d288_k768_overlap_b1_n2560_k768_d288", + "params": { + "B": 1, + "D": 288, + "K": 768, + "N": 2560, + "dtype": "bfloat16", + "seed": 8092881, + "source": "guard_overlap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.3579, + "evolution_flashlib_ms": 0.480287, + "evolution_kernel_ms": 0.18432, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.6057, + "evolution_tflops": 6.144 + } + }, + { + "label": "adjacent_c44f_d288_highk_heldout_b2_n768_k4096_d288", + "params": { + "B": 2, + "D": 288, + "K": 4096, + "N": 768, + "dtype": "bfloat16", + "seed": 4428801, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.2892, + "evolution_flashlib_ms": 1.583007, + "evolution_kernel_ms": 0.246496, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 6.422, + "evolution_tflops": 14.7016 + } + }, + { + "label": "adjacent_5600_d288_mid_highk_b2_n1024_k2048_d288", + "params": { + "B": 2, + "D": 288, + "K": 2048, + "N": 1024, + "dtype": "bfloat16", + "seed": 5602881, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.4751, + "evolution_flashlib_ms": 0.976096, + "evolution_kernel_ms": 0.206512, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 4.7266, + "evolution_tflops": 11.6987 + } + }, + { + "label": "adjacent_68cf_d288_boundary_b2_n1920_k512_d288", + "params": { + "B": 2, + "D": 288, + "K": 512, + "N": 1920, + "dtype": "bfloat16", + "seed": 6828801, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.6899, + "evolution_flashlib_ms": 0.421007, + "evolution_kernel_ms": 0.177824, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.3676, + "evolution_tflops": 6.3684 + } + }, + { + "label": "post_d895_d288_b2_n2048_k1024_d288", + "params": { + "B": 2, + "D": 288, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "seed": 228801, + "source": "new_medium_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 3.7159, + "evolution_flashlib_ms": 0.65016, + "evolution_kernel_ms": 0.18616, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.4925, + "evolution_tflops": 12.9777 + } + }, + { + "label": "adjacent_1d49_d288_low_n_heldout_b3_n1152_k2048_d288", + "params": { + "B": 3, + "D": 288, + "K": 2048, + "N": 1152, + "dtype": "bfloat16", + "seed": 1492881, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 4.0485, + "evolution_flashlib_ms": 1.007006, + "evolution_kernel_ms": 0.207744, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 4.8473, + "evolution_tflops": 19.6245 + } + }, + { + "label": "adjacent_3328_d288_guard_overlap_b4_n8192_k256_d288", + "params": { + "B": 4, + "D": 288, + "K": 256, + "N": 8192, + "dtype": "bfloat16", + "seed": 3328288, + "source": "guard_overlap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 14.007, + "evolution_flashlib_ms": 0.344959, + "evolution_kernel_ms": 0.182432, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.8909, + "evolution_tflops": 26.4857 + } + }, + { + "label": "post_d895_d288_b4_n8192_k1024_d288", + "params": { + "B": 4, + "D": 288, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "seed": 228802, + "source": "new_medium_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 28.8551, + "evolution_flashlib_ms": 0.669808, + "evolution_kernel_ms": 0.214687, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.1199, + "evolution_tflops": 90.0257 + } + }, + { + "label": "post_d895_d352_b1_n256_k256_d352", + "params": { + "B": 1, + "D": 352, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "seed": 235204, + "source": "high_small_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0736, + "evolution_flashlib_ms": 0.627167, + "evolution_kernel_ms": 0.172703, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.6315, + "evolution_tflops": 0.2671 + } + }, + { + "label": "adjacent_9ca0_d352_splitk_boundary_b1_n512_k8192_d352", + "params": { + "B": 1, + "D": 352, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 935205, + "source": "request_specific" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0601, + "evolution_flashlib_ms": 49.145822, + "evolution_kernel_ms": 0.275519, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 178.3754, + "evolution_tflops": 10.7172 + } + }, + { + "label": "post_d895_d352_b1_n512_k8192_d352", + "params": { + "B": 1, + "D": 352, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 235203, + "source": "high_k_low_n" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0616, + "evolution_flashlib_ms": 47.910581, + "evolution_kernel_ms": 0.272767, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 175.6465, + "evolution_tflops": 10.8253 + } + }, + { + "label": "adjacent_c44f_d352_random_b1_n3328_k768_d352", + "params": { + "B": 1, + "D": 352, + "K": 768, + "N": 3328, + "dtype": "bfloat16", + "seed": 4435201, + "source": "random_legal" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 3.6993, + "evolution_flashlib_ms": 0.4864, + "evolution_kernel_ms": 0.187632, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.5923, + "evolution_tflops": 9.5898 + } + }, + { + "label": "adjacent_5600_d352_splitk_edge_b2_n768_k4096_d352", + "params": { + "B": 2, + "D": 352, + "K": 4096, + "N": 768, + "dtype": "bfloat16", + "seed": 5603521, + "source": "request_specific" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.7323, + "evolution_flashlib_ms": 1.621022, + "evolution_kernel_ms": 1.051615, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.5415, + "evolution_tflops": 4.2118 + } + }, + { + "label": "adjacent_1d49_d352_request_highk_b2_n1024_k8192_d352", + "params": { + "B": 2, + "D": 352, + "K": 8192, + "N": 1024, + "dtype": "bfloat16", + "seed": 1493521, + "source": "request_specific" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 4.5385, + "evolution_flashlib_ms": 2.602461, + "evolution_kernel_ms": 0.302655, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 8.5988, + "evolution_tflops": 39.0252 + } + }, + { + "label": "post_d895_d352_b2_n2048_k1024_d352", + "params": { + "B": 2, + "D": 352, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "seed": 235201, + "source": "new_high_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 4.882, + "evolution_flashlib_ms": 0.604831, + "evolution_kernel_ms": 0.189184, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.1971, + "evolution_tflops": 15.608 + } + }, + { + "label": "adjacent_ef00_d352_request_low_n_b3_n896_k8192_d352", + "params": { + "B": 3, + "D": 352, + "K": 8192, + "N": 896, + "dtype": "bfloat16", + "seed": 9003521, + "source": "request_specific" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 4.9423, + "evolution_flashlib_ms": 3.136652, + "evolution_kernel_ms": 0.316191, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 9.9201, + "evolution_tflops": 49.0278 + } + }, + { + "label": "adjacent_3328_d352_random_legal_b3_n2048_k768_d352", + "params": { + "B": 3, + "D": 352, + "K": 768, + "N": 2048, + "dtype": "bfloat16", + "seed": 3328352, + "source": "random_legal" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 5.7044, + "evolution_flashlib_ms": 0.582335, + "evolution_kernel_ms": 0.183872, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.1671, + "evolution_tflops": 18.0663 + } + }, + { + "label": "adjacent_68cf_d352_tail_b3_n2816_k768_d352", + "params": { + "B": 3, + "D": 352, + "K": 768, + "N": 2816, + "dtype": "bfloat16", + "seed": 6835201, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 9.1314, + "evolution_flashlib_ms": 0.500207, + "evolution_kernel_ms": 0.18768, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.6652, + "evolution_tflops": 24.3372 + } + }, + { + "label": "adjacent_a2f8_d352_smallk_boundary_b4_n1024_k256_d352", + "params": { + "B": 4, + "D": 352, + "K": 256, + "N": 1024, + "dtype": "bfloat16", + "seed": 1028352, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 1.5395, + "evolution_flashlib_ms": 0.479503, + "evolution_kernel_ms": 0.176096, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.723, + "evolution_tflops": 4.192 + } + }, + { + "label": "adjacent_8f09_d352_smallk_large_n_b4_n4096_k256_d352", + "params": { + "B": 4, + "D": 352, + "K": 256, + "N": 4096, + "dtype": "bfloat16", + "seed": 8093521, + "source": "random_legal" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 9.629, + "evolution_flashlib_ms": 0.306655, + "evolution_kernel_ms": 0.176128, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.7411, + "evolution_tflops": 16.765 + } + }, + { + "label": "post_d895_d352_b4_n8192_k1024_d352", + "params": { + "B": 4, + "D": 352, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "seed": 235202, + "source": "new_high_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 36.2395, + "evolution_flashlib_ms": 0.651839, + "evolution_kernel_ms": 0.220992, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.9496, + "evolution_tflops": 106.8922 + } + }, + { + "label": "adjacent_d9d5_d352_random_b5_n2304_k768_d352", + "params": { + "B": 5, + "D": 352, + "K": 768, + "N": 2304, + "dtype": "bfloat16", + "seed": 9535201, + "source": "random_legal" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 10.6911, + "evolution_flashlib_ms": 0.582592, + "evolution_kernel_ms": 0.186944, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.1164, + "evolution_tflops": 33.3177 + } + }, + { + "label": "post_d895_d416_b1_n256_k256_d416", + "params": { + "B": 1, + "D": 416, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "seed": 241604, + "source": "high_small_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0842, + "evolution_flashlib_ms": 0.647807, + "evolution_kernel_ms": 0.175137, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.6989, + "evolution_tflops": 0.3113 + } + }, + { + "label": "adjacent_3328_d416_highk_low_n_b1_n384_k8192_d416", + "params": { + "B": 1, + "D": 416, + "K": 8192, + "N": 384, + "dtype": "bfloat16", + "seed": 3328416, + "source": "request_specific" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.063, + "evolution_flashlib_ms": 41.538691, + "evolution_kernel_ms": 0.276031, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 150.4856, + "evolution_tflops": 9.4817 + } + }, + { + "label": "post_d895_d416_b1_n512_k8192_d416", + "params": { + "B": 1, + "D": 416, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 241603, + "source": "high_k_low_n" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0852, + "evolution_flashlib_ms": 40.967689, + "evolution_kernel_ms": 0.275103, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 148.9176, + "evolution_tflops": 12.6849 + } + }, + { + "label": "adjacent_9ca0_d416_splitk_min_b1_n1024_k4096_d416", + "params": { + "B": 1, + "D": 416, + "K": 4096, + "N": 1024, + "dtype": "bfloat16", + "seed": 941605, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0848, + "evolution_flashlib_ms": 41.154087, + "evolution_kernel_ms": 0.274271, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 150.049, + "evolution_tflops": 12.7234 + } + }, + { + "label": "adjacent_ef00_d416_random_midk_b1_n2176_k768_d416", + "params": { + "B": 1, + "D": 416, + "K": 768, + "N": 2176, + "dtype": "bfloat16", + "seed": 9004161, + "source": "random_legal" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.6268, + "evolution_flashlib_ms": 0.529311, + "evolution_kernel_ms": 0.191648, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.7619, + "evolution_tflops": 7.255 + } + }, + { + "label": "adjacent_d9d5_d416_highk_request_b2_n640_k8192_d416", + "params": { + "B": 2, + "D": 416, + "K": 8192, + "N": 640, + "dtype": "bfloat16", + "seed": 9541601, + "source": "request_specific" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.6807, + "evolution_flashlib_ms": 3.254428, + "evolution_kernel_ms": 0.298432, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 10.9051, + "evolution_tflops": 29.2333 + } + }, + { + "label": "post_d895_d416_b2_n2048_k1024_d416", + "params": { + "B": 2, + "D": 416, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "seed": 241601, + "source": "new_high_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 5.4317, + "evolution_flashlib_ms": 0.642463, + "evolution_kernel_ms": 0.193983, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.312, + "evolution_tflops": 17.9895 + } + }, + { + "label": "adjacent_68cf_d416_overlap_b2_n2304_k1024_d416", + "params": { + "B": 2, + "D": 416, + "K": 1024, + "N": 2304, + "dtype": "bfloat16", + "seed": 6841601, + "source": "guard_overlap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 6.4021, + "evolution_flashlib_ms": 0.613215, + "evolution_kernel_ms": 0.196895, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.1144, + "evolution_tflops": 19.9389 + } + }, + { + "label": "adjacent_a2f8_d416_random_b2_n2560_k768_d416", + "params": { + "B": 2, + "D": 416, + "K": 768, + "N": 2560, + "dtype": "bfloat16", + "seed": 1028416, + "source": "random_legal" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 6.6441, + "evolution_flashlib_ms": 0.492399, + "evolution_kernel_ms": 0.189631, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.5966, + "evolution_tflops": 17.2522 + } + }, + { + "label": "adjacent_c44f_d416_highk_request_b3_n512_k8192_d416", + "params": { + "B": 3, + "D": 416, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 4441601, + "source": "request_specific" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.7386, + "evolution_flashlib_ms": 3.822732, + "evolution_kernel_ms": 0.303423, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 12.5987, + "evolution_tflops": 34.5029 + } + }, + { + "label": "adjacent_8f09_d416_midk_tail_b3_n3456_k768_d416", + "params": { + "B": 3, + "D": 416, + "K": 768, + "N": 3456, + "dtype": "bfloat16", + "seed": 8094161, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 13.4164, + "evolution_flashlib_ms": 0.493791, + "evolution_kernel_ms": 0.189088, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.6114, + "evolution_tflops": 35.0361 + } + }, + { + "label": "adjacent_5600_d416_smallk_boundary_b4_n2048_k256_d416", + "params": { + "B": 4, + "D": 416, + "K": 256, + "N": 2048, + "dtype": "bfloat16", + "seed": 5604161, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 4.7141, + "evolution_flashlib_ms": 0.370128, + "evolution_kernel_ms": 0.175743, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.1061, + "evolution_tflops": 9.9283 + } + }, + { + "label": "adjacent_1d49_d416_random_b4_n3840_k512_d416", + "params": { + "B": 4, + "D": 416, + "K": 512, + "N": 3840, + "dtype": "bfloat16", + "seed": 1494161, + "source": "random_legal" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 16.9546, + "evolution_flashlib_ms": 0.385919, + "evolution_kernel_ms": 0.184, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.0974, + "evolution_tflops": 35.5604 + } + }, + { + "label": "post_d895_d416_b4_n8192_k1024_d416", + "params": { + "B": 4, + "D": 416, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "seed": 241602, + "source": "new_high_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 39.3193, + "evolution_flashlib_ms": 0.710015, + "evolution_kernel_ms": 0.230463, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.0808, + "evolution_tflops": 121.1354 + } + }, + { + "label": "adjacent_3328_d480_min_boundary_b1_n128_k256_d480", + "params": { + "B": 1, + "D": 480, + "K": 256, + "N": 128, + "dtype": "bfloat16", + "seed": 3328480, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0644, + "evolution_flashlib_ms": 0.488575, + "evolution_kernel_ms": 0.177056, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.7594, + "evolution_tflops": 0.1777 + } + }, + { + "label": "post_d895_d480_b1_n256_k256_d480", + "params": { + "B": 1, + "D": 480, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "seed": 248004, + "source": "high_small_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0969, + "evolution_flashlib_ms": 0.649151, + "evolution_kernel_ms": 0.176655, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.6747, + "evolution_tflops": 0.3561 + } + }, + { + "label": "post_d895_d480_b1_n512_k8192_d480", + "params": { + "B": 1, + "D": 480, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 248003, + "source": "high_k_low_n" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.1226, + "evolution_flashlib_ms": 32.836689, + "evolution_kernel_ms": 0.278559, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 117.8803, + "evolution_tflops": 14.4548 + } + }, + { + "label": "adjacent_a2f8_d480_highk_request_b1_n896_k4096_d480", + "params": { + "B": 1, + "D": 480, + "K": 4096, + "N": 896, + "dtype": "bfloat16", + "seed": 1028480, + "source": "request_specific" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.1098, + "evolution_flashlib_ms": 32.096809, + "evolution_kernel_ms": 0.274432, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 116.9572, + "evolution_tflops": 12.8382 + } + }, + { + "label": "adjacent_5600_d480_random_b1_n1536_k1024_d480", + "params": { + "B": 1, + "D": 480, + "K": 1024, + "N": 1536, + "dtype": "bfloat16", + "seed": 5604801, + "source": "random_legal" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.1001, + "evolution_flashlib_ms": 0.718975, + "evolution_kernel_ms": 0.199488, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.6041, + "evolution_tflops": 7.5691 + } + }, + { + "label": "adjacent_8f09_d480_highk_low_n_b2_n640_k4096_d480", + "params": { + "B": 2, + "D": 480, + "K": 4096, + "N": 640, + "dtype": "bfloat16", + "seed": 8094801, + "source": "request_specific" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.5174, + "evolution_flashlib_ms": 1.999325, + "evolution_kernel_ms": 0.284192, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 7.0351, + "evolution_tflops": 17.7104 + } + }, + { + "label": "post_d895_d480_b2_n2048_k1024_d480", + "params": { + "B": 2, + "D": 480, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "seed": 248001, + "source": "new_high_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 5.9879, + "evolution_flashlib_ms": 0.672447, + "evolution_kernel_ms": 0.197919, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.3976, + "evolution_tflops": 20.3443 + } + }, + { + "label": "adjacent_1d49_d480_smallk_boundary_b2_n2816_k256_d480", + "params": { + "B": 2, + "D": 480, + "K": 256, + "N": 2816, + "dtype": "bfloat16", + "seed": 1494801, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 4.1626, + "evolution_flashlib_ms": 0.332511, + "evolution_kernel_ms": 0.178015, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.8679, + "evolution_tflops": 7.7753 + } + }, + { + "label": "adjacent_9ca0_d480_splitd_large_n_b2_n4096_k512_d480", + "params": { + "B": 2, + "D": 480, + "K": 512, + "N": 4096, + "dtype": "bfloat16", + "seed": 948005, + "source": "random_legal" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 8.9917, + "evolution_flashlib_ms": 0.447807, + "evolution_kernel_ms": 0.185087, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.4194, + "evolution_tflops": 21.7548 + } + }, + { + "label": "adjacent_d9d5_d480_heldout_highk_b3_n1024_k4096_d480", + "params": { + "B": 3, + "D": 480, + "K": 4096, + "N": 1024, + "dtype": "bfloat16", + "seed": 9548001, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 6.283, + "evolution_flashlib_ms": 1.92259, + "evolution_kernel_ms": 0.297632, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 6.4596, + "evolution_tflops": 40.5857 + } + }, + { + "label": "adjacent_ef00_d480_lowk_boundary_b3_n3200_k256_d480", + "params": { + "B": 3, + "D": 480, + "K": 256, + "N": 3200, + "dtype": "bfloat16", + "seed": 9004801, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 7.059, + "evolution_flashlib_ms": 0.334224, + "evolution_kernel_ms": 0.18192, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.8372, + "evolution_tflops": 12.9689 + } + }, + { + "label": "adjacent_68cf_d480_boundary_b4_n1664_k512_d480", + "params": { + "B": 4, + "D": 480, + "K": 512, + "N": 1664, + "dtype": "bfloat16", + "seed": 6848001, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 6.0602, + "evolution_flashlib_ms": 0.539839, + "evolution_kernel_ms": 0.18536, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.9124, + "evolution_tflops": 17.6497 + } + }, + { + "label": "post_d895_d480_b4_n8192_k1024_d480", + "params": { + "B": 4, + "D": 480, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "seed": 248002, + "source": "new_high_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 43.8869, + "evolution_flashlib_ms": 0.733983, + "evolution_kernel_ms": 0.24, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.0583, + "evolution_tflops": 134.2177 + } + }, + { + "label": "adjacent_c44f_d480_boundary_b5_n2048_k512_d480", + "params": { + "B": 5, + "D": 480, + "K": 512, + "N": 2048, + "dtype": "bfloat16", + "seed": 4448001, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 10.031, + "evolution_flashlib_ms": 0.501759, + "evolution_kernel_ms": 0.185136, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.7102, + "evolution_tflops": 27.1863 + } + } +] diff --git a/cake_exports/kmeans/benchmarks/workload.py b/cake_exports/kmeans/benchmarks/workload.py new file mode 100644 index 00000000..dff45603 --- /dev/null +++ b/cake_exports/kmeans/benchmarks/workload.py @@ -0,0 +1,29 @@ +"""Workload adapter for semantic correctness and per-shape performance. + +Replace this template or pass ``--benchmark-adapter`` to the Cake exporter. +The exported repository never needs Weave IR: this module calls its public +Python API and provides an independent reference implementation. +""" + +from __future__ import annotations + +from typing import Any + +CONFIGURED = False +SHAPES: list[dict[str, Any]] = [] + + +def make_case(package: Any, shape: dict[str, Any]) -> dict[str, Any]: + """Return run/reference/compare callables and optional work estimates. + + Required keys in the returned mapping: + run: zero-argument callable invoking the exported semantic API + reference: zero-argument callable computing independent expected output + compare: callable(actual, expected) returning bool or {"passed": bool, ...} + + Optional keys: + flops: useful floating-point operations for TFLOPS reporting + bytes: useful bytes moved for GB/s reporting + metrics: static metadata copied into the result + """ + raise NotImplementedError("configure benchmarks/workload.py before running shape benchmarks") diff --git a/cake_exports/kmeans/pyproject.toml b/cake_exports/kmeans/pyproject.toml new file mode 100644 index 00000000..331d7bb4 --- /dev/null +++ b/cake_exports/kmeans/pyproject.toml @@ -0,0 +1,40 @@ +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "kmeans" +version = "0.1.0" +requires-python = ">=3.11" +dependencies = [ + "cuda-python", + "torch", +] + +[project.optional-dependencies] +test = [ + "pytest", +] +benchmark = [ + "cupti-python", + "cuda-pathfinder", + "nvidia-cuda-cupti", + "pytest", + "triton", +] +tvm-ffi = [ + "cake-std==0.1.13.dev20260704+g7b8dbc8", +] + +[tool.pytest.ini_options] +markers = [ + "export_validation_shape: one declared contract shape counted by the fail-closed publication gate", + "gpu: requires a CUDA GPU", +] + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.setuptools.package-data] +"flashlib_cake_kmeans" = ["*.json", "*.cu", "cuda/*.cu"] + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/__init__.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/__init__.py new file mode 100644 index 00000000..54f0018b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/__init__.py @@ -0,0 +1,1455 @@ +from .kernels import ( + KERNELS, + ExportedKernel, + get_kernel, + dispatch_kernel_0000, + dispatch_kernel_0001, + dispatch_kernel_0002, + dispatch_kernel_0003, + dispatch_kernel_0004, + dispatch_kernel_0005, + dispatch_kernel_0006, + dispatch_kernel_0007, + dispatch_kernel_0008, + dispatch_kernel_0009, + dispatch_kernel_0010, + dispatch_kernel_0011, + dispatch_kernel_0012, + dispatch_kernel_0013, + dispatch_kernel_0014, + dispatch_kernel_0015, + dispatch_kernel_0016, + dispatch_kernel_0017, + dispatch_kernel_0018, + dispatch_kernel_0019, + dispatch_kernel_0020, + dispatch_kernel_0021, + dispatch_kernel_0022, + dispatch_kernel_0023, + dispatch_kernel_0024, + dispatch_kernel_0025, + dispatch_kernel_0026, + dispatch_kernel_0027, + dispatch_kernel_0028, + dispatch_kernel_0029, + dispatch_kernel_0030, + dispatch_kernel_0031, + dispatch_kernel_0032, + dispatch_kernel_0033, + dispatch_kernel_0034, + dispatch_kernel_0035, + dispatch_kernel_0036, + dispatch_kernel_0037, + dispatch_kernel_0038, + dispatch_kernel_0039, + dispatch_kernel_0040, + dispatch_kernel_0041, + dispatch_kernel_0042, + dispatch_kernel_0043, + dispatch_kernel_0044, + dispatch_kernel_0045, + dispatch_kernel_0046, + dispatch_kernel_0047, + dispatch_kernel_0048, + dispatch_kernel_0049, + dispatch_kernel_0050, + dispatch_kernel_0051, + dispatch_kernel_0052, + dispatch_kernel_0053, + dispatch_kernel_0054, + dispatch_kernel_0055, + dispatch_kernel_0056, + dispatch_kernel_0057, + dispatch_kernel_0058, + dispatch_kernel_0059, + dispatch_kernel_0060, + dispatch_kernel_0061, + dispatch_kernel_0062, + dispatch_kernel_0063, + dispatch_kernel_0064, + dispatch_kernel_0065, + dispatch_kernel_0066, + dispatch_kernel_0067, + dispatch_kernel_0068, + dispatch_kernel_0069, + dispatch_kernel_0070, + dispatch_kernel_0071, + dispatch_kernel_0072, + dispatch_kernel_0073, + dispatch_kernel_0074, + dispatch_kernel_0075, + dispatch_kernel_0076, + dispatch_kernel_0077, + dispatch_kernel_0078, + dispatch_kernel_0079, + dispatch_kernel_0080, + dispatch_kernel_0081, + dispatch_kernel_0082, + dispatch_kernel_0083, + dispatch_kernel_0084, + dispatch_kernel_0085, + dispatch_kernel_0086, + dispatch_kernel_0087, + dispatch_kernel_0088, + dispatch_kernel_0089, + dispatch_kernel_0090, + dispatch_kernel_0091, + dispatch_kernel_0092, + dispatch_kernel_0093, + dispatch_kernel_0094, + dispatch_kernel_0095, + dispatch_kernel_0096, + dispatch_kernel_0097, + dispatch_kernel_0098, + dispatch_kernel_0099, + dispatch_kernel_0100, + dispatch_kernel_0101, + dispatch_kernel_0102, + dispatch_kernel_0103, + dispatch_kernel_0104, + dispatch_kernel_0105, + dispatch_kernel_0106, + dispatch_kernel_0107, + dispatch_kernel_0108, + dispatch_kernel_0109, + dispatch_kernel_0110, + dispatch_kernel_0111, + dispatch_kernel_0112, + dispatch_kernel_0113, + dispatch_kernel_0114, + dispatch_kernel_0115, + dispatch_kernel_0116, + dispatch_kernel_0117, + dispatch_kernel_0118, + dispatch_kernel_0119, + dispatch_kernel_0120, + dispatch_kernel_0121, + dispatch_kernel_0122, + dispatch_kernel_0123, + dispatch_kernel_0124, + dispatch_kernel_0125, + dispatch_kernel_0126, + dispatch_kernel_0127, + dispatch_kernel_0128, + dispatch_kernel_0129, + dispatch_kernel_0130, + dispatch_kernel_0131, + dispatch_kernel_0132, + dispatch_kernel_0133, + dispatch_kernel_0134, + dispatch_kernel_0135, + dispatch_kernel_0136, + dispatch_kernel_0137, + dispatch_kernel_0138, + dispatch_kernel_0139, + dispatch_kernel_0140, + dispatch_kernel_0141, + dispatch_kernel_0142, + dispatch_kernel_0143, + dispatch_kernel_0144, + dispatch_kernel_0145, + dispatch_kernel_0146, + dispatch_kernel_0147, + dispatch_kernel_0148, + dispatch_kernel_0149, + dispatch_kernel_0150, + dispatch_kernel_0151, + dispatch_kernel_0152, + dispatch_kernel_0153, + dispatch_kernel_0154, + dispatch_kernel_0155, + dispatch_kernel_0156, + dispatch_kernel_0157, + dispatch_kernel_0158, + dispatch_kernel_0159, + dispatch_kernel_0160, + dispatch_kernel_0161, + dispatch_kernel_0162, + dispatch_kernel_0163, + dispatch_kernel_0164, + dispatch_kernel_0165, + dispatch_kernel_0166, + dispatch_kernel_0167, + dispatch_kernel_0168, + dispatch_kernel_0169, + dispatch_kernel_0170, + dispatch_kernel_0171, + dispatch_kernel_0172, + dispatch_kernel_0173, + dispatch_kernel_0174, + dispatch_kernel_0175, + dispatch_kernel_0176, + dispatch_kernel_0177, + dispatch_kernel_0178, + dispatch_kernel_0179, + dispatch_kernel_0180, + dispatch_kernel_0181, + dispatch_kernel_0182, + dispatch_kernel_0183, + dispatch_kernel_0184, + dispatch_kernel_0185, + dispatch_kernel_0186, + dispatch_kernel_0187, + dispatch_kernel_0188, + dispatch_kernel_0189, + dispatch_kernel_0190, + dispatch_kernel_0191, + dispatch_kernel_0192, + dispatch_kernel_0193, + dispatch_kernel_0194, + dispatch_kernel_0195, + dispatch_kernel_0196, + dispatch_kernel_0197, + dispatch_kernel_0198, + dispatch_kernel_0199, + dispatch_kernel_0200, + dispatch_kernel_0201, + dispatch_kernel_0202, + dispatch_kernel_0203, + dispatch_kernel_0204, + dispatch_kernel_0205, + dispatch_kernel_0206, + dispatch_kernel_0207, + dispatch_kernel_0208, + dispatch_kernel_0209, + dispatch_kernel_0210, + dispatch_kernel_0211, + dispatch_kernel_0212, + dispatch_kernel_0213, + dispatch_kernel_0214, + dispatch_kernel_0215, + dispatch_kernel_0216, + dispatch_kernel_0217, + dispatch_kernel_0218, + dispatch_kernel_0219, + dispatch_kernel_0220, + dispatch_kernel_0221, + dispatch_kernel_0222, + dispatch_kernel_0223, + dispatch_kernel_0224, + dispatch_kernel_0225, + dispatch_kernel_0226, + dispatch_kernel_0227, + dispatch_kernel_0228, + dispatch_kernel_0229, + dispatch_kernel_0230, + dispatch_kernel_0231, + dispatch_kernel_0232, + dispatch_kernel_0233, + dispatch_kernel_0234, + dispatch_kernel_0235, + dispatch_kernel_0236, + dispatch_kernel_0237, + dispatch_kernel_0238, + dispatch_kernel_0239, + dispatch_kernel_0240, + dispatch_kernel_0241, + dispatch_kernel_0242, + dispatch_kernel_0243, + dispatch_kernel_0244, + dispatch_kernel_0245, + dispatch_kernel_0246, + dispatch_kernel_0247, + dispatch_kernel_0248, + dispatch_kernel_0249, + dispatch_kernel_0250, + dispatch_kernel_0251, + dispatch_kernel_0252, + dispatch_kernel_0253, + dispatch_kernel_0254, + dispatch_kernel_0255, + dispatch_kernel_0256, + dispatch_kernel_0257, + dispatch_kernel_0258, + dispatch_kernel_0259, + dispatch_kernel_0260, + dispatch_kernel_0261, + dispatch_kernel_0262, + dispatch_kernel_0263, + dispatch_kernel_0264, + dispatch_kernel_0265, + dispatch_kernel_0266, + dispatch_kernel_0267, + dispatch_kernel_0268, + dispatch_kernel_0269, + dispatch_kernel_0270, + dispatch_kernel_0271, + dispatch_kernel_0272, + dispatch_kernel_0273, + dispatch_kernel_0274, + dispatch_kernel_0275, + dispatch_kernel_0276, + dispatch_kernel_0277, + dispatch_kernel_0278, + dispatch_kernel_0279, + dispatch_kernel_0280, + dispatch_kernel_0281, + dispatch_kernel_0282, + dispatch_kernel_0283, + dispatch_kernel_0284, + dispatch_kernel_0285, + dispatch_kernel_0286, + dispatch_kernel_0287, + dispatch_kernel_0288, + dispatch_kernel_0289, + dispatch_kernel_0290, + dispatch_kernel_0291, + dispatch_kernel_0292, + dispatch_kernel_0293, + dispatch_kernel_0294, + dispatch_kernel_0295, + dispatch_kernel_0296, + dispatch_kernel_0297, + dispatch_kernel_0298, + dispatch_kernel_0299, + dispatch_kernel_0300, + dispatch_kernel_0301, + dispatch_kernel_0302, + dispatch_kernel_0303, + dispatch_kernel_0304, + dispatch_kernel_0305, + dispatch_kernel_0306, + dispatch_kernel_0307, + dispatch_kernel_0308, + dispatch_kernel_0309, + dispatch_kernel_0310, + dispatch_kernel_0311, + dispatch_kernel_0312, + dispatch_kernel_0313, + dispatch_kernel_0314, + dispatch_kernel_0315, + dispatch_kernel_0316, + dispatch_kernel_0317, + dispatch_kernel_0318, + dispatch_kernel_0319, + dispatch_kernel_0320, + dispatch_kernel_0321, + dispatch_kernel_0322, + dispatch_kernel_0323, + dispatch_kernel_0324, + dispatch_kernel_0325, + dispatch_kernel_0326, + dispatch_kernel_0327, + dispatch_kernel_0328, + dispatch_kernel_0329, + dispatch_kernel_0330, + dispatch_kernel_0331, + dispatch_kernel_0332, + dispatch_kernel_0333, + dispatch_kernel_0334, + dispatch_kernel_0335, + dispatch_kernel_0336, + dispatch_kernel_0337, + dispatch_kernel_0338, + dispatch_kernel_0339, + dispatch_kernel_0340, + dispatch_kernel_0341, + dispatch_kernel_0342, + dispatch_kernel_0343, + dispatch_kernel_0344, + dispatch_kernel_0345, + dispatch_kernel_0346, + dispatch_kernel_0347, + dispatch_kernel_0348, + dispatch_kernel_0349, + dispatch_kernel_0350, + dispatch_kernel_0351, + dispatch_kernel_0352, + dispatch_kernel_0353, + dispatch_kernel_0354, + dispatch_kernel_0355, + dispatch_kernel_0356, + dispatch_kernel_0357, + launch_dispatch_kernel_0000, + launch_dispatch_kernel_0001, + launch_dispatch_kernel_0002, + launch_dispatch_kernel_0003, + launch_dispatch_kernel_0004, + launch_dispatch_kernel_0005, + launch_dispatch_kernel_0006, + launch_dispatch_kernel_0007, + launch_dispatch_kernel_0008, + launch_dispatch_kernel_0009, + launch_dispatch_kernel_0010, + launch_dispatch_kernel_0011, + launch_dispatch_kernel_0012, + launch_dispatch_kernel_0013, + launch_dispatch_kernel_0014, + launch_dispatch_kernel_0015, + launch_dispatch_kernel_0016, + launch_dispatch_kernel_0017, + launch_dispatch_kernel_0018, + launch_dispatch_kernel_0019, + launch_dispatch_kernel_0020, + launch_dispatch_kernel_0021, + launch_dispatch_kernel_0022, + launch_dispatch_kernel_0023, + launch_dispatch_kernel_0024, + launch_dispatch_kernel_0025, + launch_dispatch_kernel_0026, + launch_dispatch_kernel_0027, + launch_dispatch_kernel_0028, + launch_dispatch_kernel_0029, + launch_dispatch_kernel_0030, + launch_dispatch_kernel_0031, + launch_dispatch_kernel_0032, + launch_dispatch_kernel_0033, + launch_dispatch_kernel_0034, + launch_dispatch_kernel_0035, + launch_dispatch_kernel_0036, + launch_dispatch_kernel_0037, + launch_dispatch_kernel_0038, + launch_dispatch_kernel_0039, + launch_dispatch_kernel_0040, + launch_dispatch_kernel_0041, + launch_dispatch_kernel_0042, + launch_dispatch_kernel_0043, + launch_dispatch_kernel_0044, + launch_dispatch_kernel_0045, + launch_dispatch_kernel_0046, + launch_dispatch_kernel_0047, + launch_dispatch_kernel_0048, + launch_dispatch_kernel_0049, + launch_dispatch_kernel_0050, + launch_dispatch_kernel_0051, + launch_dispatch_kernel_0052, + launch_dispatch_kernel_0053, + launch_dispatch_kernel_0054, + launch_dispatch_kernel_0055, + launch_dispatch_kernel_0056, + launch_dispatch_kernel_0057, + launch_dispatch_kernel_0058, + launch_dispatch_kernel_0059, + launch_dispatch_kernel_0060, + launch_dispatch_kernel_0061, + launch_dispatch_kernel_0062, + launch_dispatch_kernel_0063, + launch_dispatch_kernel_0064, + launch_dispatch_kernel_0065, + launch_dispatch_kernel_0066, + launch_dispatch_kernel_0067, + launch_dispatch_kernel_0068, + launch_dispatch_kernel_0069, + launch_dispatch_kernel_0070, + launch_dispatch_kernel_0071, + launch_dispatch_kernel_0072, + launch_dispatch_kernel_0073, + launch_dispatch_kernel_0074, + launch_dispatch_kernel_0075, + launch_dispatch_kernel_0076, + launch_dispatch_kernel_0077, + launch_dispatch_kernel_0078, + launch_dispatch_kernel_0079, + launch_dispatch_kernel_0080, + launch_dispatch_kernel_0081, + launch_dispatch_kernel_0082, + launch_dispatch_kernel_0083, + launch_dispatch_kernel_0084, + launch_dispatch_kernel_0085, + launch_dispatch_kernel_0086, + launch_dispatch_kernel_0087, + launch_dispatch_kernel_0088, + launch_dispatch_kernel_0089, + launch_dispatch_kernel_0090, + launch_dispatch_kernel_0091, + launch_dispatch_kernel_0092, + launch_dispatch_kernel_0093, + launch_dispatch_kernel_0094, + launch_dispatch_kernel_0095, + launch_dispatch_kernel_0096, + launch_dispatch_kernel_0097, + launch_dispatch_kernel_0098, + launch_dispatch_kernel_0099, + launch_dispatch_kernel_0100, + launch_dispatch_kernel_0101, + launch_dispatch_kernel_0102, + launch_dispatch_kernel_0103, + launch_dispatch_kernel_0104, + launch_dispatch_kernel_0105, + launch_dispatch_kernel_0106, + launch_dispatch_kernel_0107, + launch_dispatch_kernel_0108, + launch_dispatch_kernel_0109, + launch_dispatch_kernel_0110, + launch_dispatch_kernel_0111, + launch_dispatch_kernel_0112, + launch_dispatch_kernel_0113, + launch_dispatch_kernel_0114, + launch_dispatch_kernel_0115, + launch_dispatch_kernel_0116, + launch_dispatch_kernel_0117, + launch_dispatch_kernel_0118, + launch_dispatch_kernel_0119, + launch_dispatch_kernel_0120, + launch_dispatch_kernel_0121, + launch_dispatch_kernel_0122, + launch_dispatch_kernel_0123, + launch_dispatch_kernel_0124, + launch_dispatch_kernel_0125, + launch_dispatch_kernel_0126, + launch_dispatch_kernel_0127, + launch_dispatch_kernel_0128, + launch_dispatch_kernel_0129, + launch_dispatch_kernel_0130, + launch_dispatch_kernel_0131, + launch_dispatch_kernel_0132, + launch_dispatch_kernel_0133, + launch_dispatch_kernel_0134, + launch_dispatch_kernel_0135, + launch_dispatch_kernel_0136, + launch_dispatch_kernel_0137, + launch_dispatch_kernel_0138, + launch_dispatch_kernel_0139, + launch_dispatch_kernel_0140, + launch_dispatch_kernel_0141, + launch_dispatch_kernel_0142, + launch_dispatch_kernel_0143, + launch_dispatch_kernel_0144, + launch_dispatch_kernel_0145, + launch_dispatch_kernel_0146, + launch_dispatch_kernel_0147, + launch_dispatch_kernel_0148, + launch_dispatch_kernel_0149, + launch_dispatch_kernel_0150, + launch_dispatch_kernel_0151, + launch_dispatch_kernel_0152, + launch_dispatch_kernel_0153, + launch_dispatch_kernel_0154, + launch_dispatch_kernel_0155, + launch_dispatch_kernel_0156, + launch_dispatch_kernel_0157, + launch_dispatch_kernel_0158, + launch_dispatch_kernel_0159, + launch_dispatch_kernel_0160, + launch_dispatch_kernel_0161, + launch_dispatch_kernel_0162, + launch_dispatch_kernel_0163, + launch_dispatch_kernel_0164, + launch_dispatch_kernel_0165, + launch_dispatch_kernel_0166, + launch_dispatch_kernel_0167, + launch_dispatch_kernel_0168, + launch_dispatch_kernel_0169, + launch_dispatch_kernel_0170, + launch_dispatch_kernel_0171, + launch_dispatch_kernel_0172, + launch_dispatch_kernel_0173, + launch_dispatch_kernel_0174, + launch_dispatch_kernel_0175, + launch_dispatch_kernel_0176, + launch_dispatch_kernel_0177, + launch_dispatch_kernel_0178, + launch_dispatch_kernel_0179, + launch_dispatch_kernel_0180, + launch_dispatch_kernel_0181, + launch_dispatch_kernel_0182, + launch_dispatch_kernel_0183, + launch_dispatch_kernel_0184, + launch_dispatch_kernel_0185, + launch_dispatch_kernel_0186, + launch_dispatch_kernel_0187, + launch_dispatch_kernel_0188, + launch_dispatch_kernel_0189, + launch_dispatch_kernel_0190, + launch_dispatch_kernel_0191, + launch_dispatch_kernel_0192, + launch_dispatch_kernel_0193, + launch_dispatch_kernel_0194, + launch_dispatch_kernel_0195, + launch_dispatch_kernel_0196, + launch_dispatch_kernel_0197, + launch_dispatch_kernel_0198, + launch_dispatch_kernel_0199, + launch_dispatch_kernel_0200, + launch_dispatch_kernel_0201, + launch_dispatch_kernel_0202, + launch_dispatch_kernel_0203, + launch_dispatch_kernel_0204, + launch_dispatch_kernel_0205, + launch_dispatch_kernel_0206, + launch_dispatch_kernel_0207, + launch_dispatch_kernel_0208, + launch_dispatch_kernel_0209, + launch_dispatch_kernel_0210, + launch_dispatch_kernel_0211, + launch_dispatch_kernel_0212, + launch_dispatch_kernel_0213, + launch_dispatch_kernel_0214, + launch_dispatch_kernel_0215, + launch_dispatch_kernel_0216, + launch_dispatch_kernel_0217, + launch_dispatch_kernel_0218, + launch_dispatch_kernel_0219, + launch_dispatch_kernel_0220, + launch_dispatch_kernel_0221, + launch_dispatch_kernel_0222, + launch_dispatch_kernel_0223, + launch_dispatch_kernel_0224, + launch_dispatch_kernel_0225, + launch_dispatch_kernel_0226, + launch_dispatch_kernel_0227, + launch_dispatch_kernel_0228, + launch_dispatch_kernel_0229, + launch_dispatch_kernel_0230, + launch_dispatch_kernel_0231, + launch_dispatch_kernel_0232, + launch_dispatch_kernel_0233, + launch_dispatch_kernel_0234, + launch_dispatch_kernel_0235, + launch_dispatch_kernel_0236, + launch_dispatch_kernel_0237, + launch_dispatch_kernel_0238, + launch_dispatch_kernel_0239, + launch_dispatch_kernel_0240, + launch_dispatch_kernel_0241, + launch_dispatch_kernel_0242, + launch_dispatch_kernel_0243, + launch_dispatch_kernel_0244, + launch_dispatch_kernel_0245, + launch_dispatch_kernel_0246, + launch_dispatch_kernel_0247, + launch_dispatch_kernel_0248, + launch_dispatch_kernel_0249, + launch_dispatch_kernel_0250, + launch_dispatch_kernel_0251, + launch_dispatch_kernel_0252, + launch_dispatch_kernel_0253, + launch_dispatch_kernel_0254, + launch_dispatch_kernel_0255, + launch_dispatch_kernel_0256, + launch_dispatch_kernel_0257, + launch_dispatch_kernel_0258, + launch_dispatch_kernel_0259, + launch_dispatch_kernel_0260, + launch_dispatch_kernel_0261, + launch_dispatch_kernel_0262, + launch_dispatch_kernel_0263, + launch_dispatch_kernel_0264, + launch_dispatch_kernel_0265, + launch_dispatch_kernel_0266, + launch_dispatch_kernel_0267, + launch_dispatch_kernel_0268, + launch_dispatch_kernel_0269, + launch_dispatch_kernel_0270, + launch_dispatch_kernel_0271, + launch_dispatch_kernel_0272, + launch_dispatch_kernel_0273, + launch_dispatch_kernel_0274, + launch_dispatch_kernel_0275, + launch_dispatch_kernel_0276, + launch_dispatch_kernel_0277, + launch_dispatch_kernel_0278, + launch_dispatch_kernel_0279, + launch_dispatch_kernel_0280, + launch_dispatch_kernel_0281, + launch_dispatch_kernel_0282, + launch_dispatch_kernel_0283, + launch_dispatch_kernel_0284, + launch_dispatch_kernel_0285, + launch_dispatch_kernel_0286, + launch_dispatch_kernel_0287, + launch_dispatch_kernel_0288, + launch_dispatch_kernel_0289, + launch_dispatch_kernel_0290, + launch_dispatch_kernel_0291, + launch_dispatch_kernel_0292, + launch_dispatch_kernel_0293, + launch_dispatch_kernel_0294, + launch_dispatch_kernel_0295, + launch_dispatch_kernel_0296, + launch_dispatch_kernel_0297, + launch_dispatch_kernel_0298, + launch_dispatch_kernel_0299, + launch_dispatch_kernel_0300, + launch_dispatch_kernel_0301, + launch_dispatch_kernel_0302, + launch_dispatch_kernel_0303, + launch_dispatch_kernel_0304, + launch_dispatch_kernel_0305, + launch_dispatch_kernel_0306, + launch_dispatch_kernel_0307, + launch_dispatch_kernel_0308, + launch_dispatch_kernel_0309, + launch_dispatch_kernel_0310, + launch_dispatch_kernel_0311, + launch_dispatch_kernel_0312, + launch_dispatch_kernel_0313, + launch_dispatch_kernel_0314, + launch_dispatch_kernel_0315, + launch_dispatch_kernel_0316, + launch_dispatch_kernel_0317, + launch_dispatch_kernel_0318, + launch_dispatch_kernel_0319, + launch_dispatch_kernel_0320, + launch_dispatch_kernel_0321, + launch_dispatch_kernel_0322, + launch_dispatch_kernel_0323, + launch_dispatch_kernel_0324, + launch_dispatch_kernel_0325, + launch_dispatch_kernel_0326, + launch_dispatch_kernel_0327, + launch_dispatch_kernel_0328, + launch_dispatch_kernel_0329, + launch_dispatch_kernel_0330, + launch_dispatch_kernel_0331, + launch_dispatch_kernel_0332, + launch_dispatch_kernel_0333, + launch_dispatch_kernel_0334, + launch_dispatch_kernel_0335, + launch_dispatch_kernel_0336, + launch_dispatch_kernel_0337, + launch_dispatch_kernel_0338, + launch_dispatch_kernel_0339, + launch_dispatch_kernel_0340, + launch_dispatch_kernel_0341, + launch_dispatch_kernel_0342, + launch_dispatch_kernel_0343, + launch_dispatch_kernel_0344, + launch_dispatch_kernel_0345, + launch_dispatch_kernel_0346, + launch_dispatch_kernel_0347, + launch_dispatch_kernel_0348, + launch_dispatch_kernel_0349, + launch_dispatch_kernel_0350, + launch_dispatch_kernel_0351, + launch_dispatch_kernel_0352, + launch_dispatch_kernel_0353, + launch_dispatch_kernel_0354, + launch_dispatch_kernel_0355, + launch_dispatch_kernel_0356, + launch_dispatch_kernel_0357, +) +from .tvm_ffi import register_tvm_ffi, tvm_ffi_function_names + +__all__ = [ + 'KERNELS', + 'ExportedKernel', + 'get_kernel', + 'dispatch_kernel_0000', + 'dispatch_kernel_0001', + 'dispatch_kernel_0002', + 'dispatch_kernel_0003', + 'dispatch_kernel_0004', + 'dispatch_kernel_0005', + 'dispatch_kernel_0006', + 'dispatch_kernel_0007', + 'dispatch_kernel_0008', + 'dispatch_kernel_0009', + 'dispatch_kernel_0010', + 'dispatch_kernel_0011', + 'dispatch_kernel_0012', + 'dispatch_kernel_0013', + 'dispatch_kernel_0014', + 'dispatch_kernel_0015', + 'dispatch_kernel_0016', + 'dispatch_kernel_0017', + 'dispatch_kernel_0018', + 'dispatch_kernel_0019', + 'dispatch_kernel_0020', + 'dispatch_kernel_0021', + 'dispatch_kernel_0022', + 'dispatch_kernel_0023', + 'dispatch_kernel_0024', + 'dispatch_kernel_0025', + 'dispatch_kernel_0026', + 'dispatch_kernel_0027', + 'dispatch_kernel_0028', + 'dispatch_kernel_0029', + 'dispatch_kernel_0030', + 'dispatch_kernel_0031', + 'dispatch_kernel_0032', + 'dispatch_kernel_0033', + 'dispatch_kernel_0034', + 'dispatch_kernel_0035', + 'dispatch_kernel_0036', + 'dispatch_kernel_0037', + 'dispatch_kernel_0038', + 'dispatch_kernel_0039', + 'dispatch_kernel_0040', + 'dispatch_kernel_0041', + 'dispatch_kernel_0042', + 'dispatch_kernel_0043', + 'dispatch_kernel_0044', + 'dispatch_kernel_0045', + 'dispatch_kernel_0046', + 'dispatch_kernel_0047', + 'dispatch_kernel_0048', + 'dispatch_kernel_0049', + 'dispatch_kernel_0050', + 'dispatch_kernel_0051', + 'dispatch_kernel_0052', + 'dispatch_kernel_0053', + 'dispatch_kernel_0054', + 'dispatch_kernel_0055', + 'dispatch_kernel_0056', + 'dispatch_kernel_0057', + 'dispatch_kernel_0058', + 'dispatch_kernel_0059', + 'dispatch_kernel_0060', + 'dispatch_kernel_0061', + 'dispatch_kernel_0062', + 'dispatch_kernel_0063', + 'dispatch_kernel_0064', + 'dispatch_kernel_0065', + 'dispatch_kernel_0066', + 'dispatch_kernel_0067', + 'dispatch_kernel_0068', + 'dispatch_kernel_0069', + 'dispatch_kernel_0070', + 'dispatch_kernel_0071', + 'dispatch_kernel_0072', + 'dispatch_kernel_0073', + 'dispatch_kernel_0074', + 'dispatch_kernel_0075', + 'dispatch_kernel_0076', + 'dispatch_kernel_0077', + 'dispatch_kernel_0078', + 'dispatch_kernel_0079', + 'dispatch_kernel_0080', + 'dispatch_kernel_0081', + 'dispatch_kernel_0082', + 'dispatch_kernel_0083', + 'dispatch_kernel_0084', + 'dispatch_kernel_0085', + 'dispatch_kernel_0086', + 'dispatch_kernel_0087', + 'dispatch_kernel_0088', + 'dispatch_kernel_0089', + 'dispatch_kernel_0090', + 'dispatch_kernel_0091', + 'dispatch_kernel_0092', + 'dispatch_kernel_0093', + 'dispatch_kernel_0094', + 'dispatch_kernel_0095', + 'dispatch_kernel_0096', + 'dispatch_kernel_0097', + 'dispatch_kernel_0098', + 'dispatch_kernel_0099', + 'dispatch_kernel_0100', + 'dispatch_kernel_0101', + 'dispatch_kernel_0102', + 'dispatch_kernel_0103', + 'dispatch_kernel_0104', + 'dispatch_kernel_0105', + 'dispatch_kernel_0106', + 'dispatch_kernel_0107', + 'dispatch_kernel_0108', + 'dispatch_kernel_0109', + 'dispatch_kernel_0110', + 'dispatch_kernel_0111', + 'dispatch_kernel_0112', + 'dispatch_kernel_0113', + 'dispatch_kernel_0114', + 'dispatch_kernel_0115', + 'dispatch_kernel_0116', + 'dispatch_kernel_0117', + 'dispatch_kernel_0118', + 'dispatch_kernel_0119', + 'dispatch_kernel_0120', + 'dispatch_kernel_0121', + 'dispatch_kernel_0122', + 'dispatch_kernel_0123', + 'dispatch_kernel_0124', + 'dispatch_kernel_0125', + 'dispatch_kernel_0126', + 'dispatch_kernel_0127', + 'dispatch_kernel_0128', + 'dispatch_kernel_0129', + 'dispatch_kernel_0130', + 'dispatch_kernel_0131', + 'dispatch_kernel_0132', + 'dispatch_kernel_0133', + 'dispatch_kernel_0134', + 'dispatch_kernel_0135', + 'dispatch_kernel_0136', + 'dispatch_kernel_0137', + 'dispatch_kernel_0138', + 'dispatch_kernel_0139', + 'dispatch_kernel_0140', + 'dispatch_kernel_0141', + 'dispatch_kernel_0142', + 'dispatch_kernel_0143', + 'dispatch_kernel_0144', + 'dispatch_kernel_0145', + 'dispatch_kernel_0146', + 'dispatch_kernel_0147', + 'dispatch_kernel_0148', + 'dispatch_kernel_0149', + 'dispatch_kernel_0150', + 'dispatch_kernel_0151', + 'dispatch_kernel_0152', + 'dispatch_kernel_0153', + 'dispatch_kernel_0154', + 'dispatch_kernel_0155', + 'dispatch_kernel_0156', + 'dispatch_kernel_0157', + 'dispatch_kernel_0158', + 'dispatch_kernel_0159', + 'dispatch_kernel_0160', + 'dispatch_kernel_0161', + 'dispatch_kernel_0162', + 'dispatch_kernel_0163', + 'dispatch_kernel_0164', + 'dispatch_kernel_0165', + 'dispatch_kernel_0166', + 'dispatch_kernel_0167', + 'dispatch_kernel_0168', + 'dispatch_kernel_0169', + 'dispatch_kernel_0170', + 'dispatch_kernel_0171', + 'dispatch_kernel_0172', + 'dispatch_kernel_0173', + 'dispatch_kernel_0174', + 'dispatch_kernel_0175', + 'dispatch_kernel_0176', + 'dispatch_kernel_0177', + 'dispatch_kernel_0178', + 'dispatch_kernel_0179', + 'dispatch_kernel_0180', + 'dispatch_kernel_0181', + 'dispatch_kernel_0182', + 'dispatch_kernel_0183', + 'dispatch_kernel_0184', + 'dispatch_kernel_0185', + 'dispatch_kernel_0186', + 'dispatch_kernel_0187', + 'dispatch_kernel_0188', + 'dispatch_kernel_0189', + 'dispatch_kernel_0190', + 'dispatch_kernel_0191', + 'dispatch_kernel_0192', + 'dispatch_kernel_0193', + 'dispatch_kernel_0194', + 'dispatch_kernel_0195', + 'dispatch_kernel_0196', + 'dispatch_kernel_0197', + 'dispatch_kernel_0198', + 'dispatch_kernel_0199', + 'dispatch_kernel_0200', + 'dispatch_kernel_0201', + 'dispatch_kernel_0202', + 'dispatch_kernel_0203', + 'dispatch_kernel_0204', + 'dispatch_kernel_0205', + 'dispatch_kernel_0206', + 'dispatch_kernel_0207', + 'dispatch_kernel_0208', + 'dispatch_kernel_0209', + 'dispatch_kernel_0210', + 'dispatch_kernel_0211', + 'dispatch_kernel_0212', + 'dispatch_kernel_0213', + 'dispatch_kernel_0214', + 'dispatch_kernel_0215', + 'dispatch_kernel_0216', + 'dispatch_kernel_0217', + 'dispatch_kernel_0218', + 'dispatch_kernel_0219', + 'dispatch_kernel_0220', + 'dispatch_kernel_0221', + 'dispatch_kernel_0222', + 'dispatch_kernel_0223', + 'dispatch_kernel_0224', + 'dispatch_kernel_0225', + 'dispatch_kernel_0226', + 'dispatch_kernel_0227', + 'dispatch_kernel_0228', + 'dispatch_kernel_0229', + 'dispatch_kernel_0230', + 'dispatch_kernel_0231', + 'dispatch_kernel_0232', + 'dispatch_kernel_0233', + 'dispatch_kernel_0234', + 'dispatch_kernel_0235', + 'dispatch_kernel_0236', + 'dispatch_kernel_0237', + 'dispatch_kernel_0238', + 'dispatch_kernel_0239', + 'dispatch_kernel_0240', + 'dispatch_kernel_0241', + 'dispatch_kernel_0242', + 'dispatch_kernel_0243', + 'dispatch_kernel_0244', + 'dispatch_kernel_0245', + 'dispatch_kernel_0246', + 'dispatch_kernel_0247', + 'dispatch_kernel_0248', + 'dispatch_kernel_0249', + 'dispatch_kernel_0250', + 'dispatch_kernel_0251', + 'dispatch_kernel_0252', + 'dispatch_kernel_0253', + 'dispatch_kernel_0254', + 'dispatch_kernel_0255', + 'dispatch_kernel_0256', + 'dispatch_kernel_0257', + 'dispatch_kernel_0258', + 'dispatch_kernel_0259', + 'dispatch_kernel_0260', + 'dispatch_kernel_0261', + 'dispatch_kernel_0262', + 'dispatch_kernel_0263', + 'dispatch_kernel_0264', + 'dispatch_kernel_0265', + 'dispatch_kernel_0266', + 'dispatch_kernel_0267', + 'dispatch_kernel_0268', + 'dispatch_kernel_0269', + 'dispatch_kernel_0270', + 'dispatch_kernel_0271', + 'dispatch_kernel_0272', + 'dispatch_kernel_0273', + 'dispatch_kernel_0274', + 'dispatch_kernel_0275', + 'dispatch_kernel_0276', + 'dispatch_kernel_0277', + 'dispatch_kernel_0278', + 'dispatch_kernel_0279', + 'dispatch_kernel_0280', + 'dispatch_kernel_0281', + 'dispatch_kernel_0282', + 'dispatch_kernel_0283', + 'dispatch_kernel_0284', + 'dispatch_kernel_0285', + 'dispatch_kernel_0286', + 'dispatch_kernel_0287', + 'dispatch_kernel_0288', + 'dispatch_kernel_0289', + 'dispatch_kernel_0290', + 'dispatch_kernel_0291', + 'dispatch_kernel_0292', + 'dispatch_kernel_0293', + 'dispatch_kernel_0294', + 'dispatch_kernel_0295', + 'dispatch_kernel_0296', + 'dispatch_kernel_0297', + 'dispatch_kernel_0298', + 'dispatch_kernel_0299', + 'dispatch_kernel_0300', + 'dispatch_kernel_0301', + 'dispatch_kernel_0302', + 'dispatch_kernel_0303', + 'dispatch_kernel_0304', + 'dispatch_kernel_0305', + 'dispatch_kernel_0306', + 'dispatch_kernel_0307', + 'dispatch_kernel_0308', + 'dispatch_kernel_0309', + 'dispatch_kernel_0310', + 'dispatch_kernel_0311', + 'dispatch_kernel_0312', + 'dispatch_kernel_0313', + 'dispatch_kernel_0314', + 'dispatch_kernel_0315', + 'dispatch_kernel_0316', + 'dispatch_kernel_0317', + 'dispatch_kernel_0318', + 'dispatch_kernel_0319', + 'dispatch_kernel_0320', + 'dispatch_kernel_0321', + 'dispatch_kernel_0322', + 'dispatch_kernel_0323', + 'dispatch_kernel_0324', + 'dispatch_kernel_0325', + 'dispatch_kernel_0326', + 'dispatch_kernel_0327', + 'dispatch_kernel_0328', + 'dispatch_kernel_0329', + 'dispatch_kernel_0330', + 'dispatch_kernel_0331', + 'dispatch_kernel_0332', + 'dispatch_kernel_0333', + 'dispatch_kernel_0334', + 'dispatch_kernel_0335', + 'dispatch_kernel_0336', + 'dispatch_kernel_0337', + 'dispatch_kernel_0338', + 'dispatch_kernel_0339', + 'dispatch_kernel_0340', + 'dispatch_kernel_0341', + 'dispatch_kernel_0342', + 'dispatch_kernel_0343', + 'dispatch_kernel_0344', + 'dispatch_kernel_0345', + 'dispatch_kernel_0346', + 'dispatch_kernel_0347', + 'dispatch_kernel_0348', + 'dispatch_kernel_0349', + 'dispatch_kernel_0350', + 'dispatch_kernel_0351', + 'dispatch_kernel_0352', + 'dispatch_kernel_0353', + 'dispatch_kernel_0354', + 'dispatch_kernel_0355', + 'dispatch_kernel_0356', + 'dispatch_kernel_0357', + 'launch_dispatch_kernel_0000', + 'launch_dispatch_kernel_0001', + 'launch_dispatch_kernel_0002', + 'launch_dispatch_kernel_0003', + 'launch_dispatch_kernel_0004', + 'launch_dispatch_kernel_0005', + 'launch_dispatch_kernel_0006', + 'launch_dispatch_kernel_0007', + 'launch_dispatch_kernel_0008', + 'launch_dispatch_kernel_0009', + 'launch_dispatch_kernel_0010', + 'launch_dispatch_kernel_0011', + 'launch_dispatch_kernel_0012', + 'launch_dispatch_kernel_0013', + 'launch_dispatch_kernel_0014', + 'launch_dispatch_kernel_0015', + 'launch_dispatch_kernel_0016', + 'launch_dispatch_kernel_0017', + 'launch_dispatch_kernel_0018', + 'launch_dispatch_kernel_0019', + 'launch_dispatch_kernel_0020', + 'launch_dispatch_kernel_0021', + 'launch_dispatch_kernel_0022', + 'launch_dispatch_kernel_0023', + 'launch_dispatch_kernel_0024', + 'launch_dispatch_kernel_0025', + 'launch_dispatch_kernel_0026', + 'launch_dispatch_kernel_0027', + 'launch_dispatch_kernel_0028', + 'launch_dispatch_kernel_0029', + 'launch_dispatch_kernel_0030', + 'launch_dispatch_kernel_0031', + 'launch_dispatch_kernel_0032', + 'launch_dispatch_kernel_0033', + 'launch_dispatch_kernel_0034', + 'launch_dispatch_kernel_0035', + 'launch_dispatch_kernel_0036', + 'launch_dispatch_kernel_0037', + 'launch_dispatch_kernel_0038', + 'launch_dispatch_kernel_0039', + 'launch_dispatch_kernel_0040', + 'launch_dispatch_kernel_0041', + 'launch_dispatch_kernel_0042', + 'launch_dispatch_kernel_0043', + 'launch_dispatch_kernel_0044', + 'launch_dispatch_kernel_0045', + 'launch_dispatch_kernel_0046', + 'launch_dispatch_kernel_0047', + 'launch_dispatch_kernel_0048', + 'launch_dispatch_kernel_0049', + 'launch_dispatch_kernel_0050', + 'launch_dispatch_kernel_0051', + 'launch_dispatch_kernel_0052', + 'launch_dispatch_kernel_0053', + 'launch_dispatch_kernel_0054', + 'launch_dispatch_kernel_0055', + 'launch_dispatch_kernel_0056', + 'launch_dispatch_kernel_0057', + 'launch_dispatch_kernel_0058', + 'launch_dispatch_kernel_0059', + 'launch_dispatch_kernel_0060', + 'launch_dispatch_kernel_0061', + 'launch_dispatch_kernel_0062', + 'launch_dispatch_kernel_0063', + 'launch_dispatch_kernel_0064', + 'launch_dispatch_kernel_0065', + 'launch_dispatch_kernel_0066', + 'launch_dispatch_kernel_0067', + 'launch_dispatch_kernel_0068', + 'launch_dispatch_kernel_0069', + 'launch_dispatch_kernel_0070', + 'launch_dispatch_kernel_0071', + 'launch_dispatch_kernel_0072', + 'launch_dispatch_kernel_0073', + 'launch_dispatch_kernel_0074', + 'launch_dispatch_kernel_0075', + 'launch_dispatch_kernel_0076', + 'launch_dispatch_kernel_0077', + 'launch_dispatch_kernel_0078', + 'launch_dispatch_kernel_0079', + 'launch_dispatch_kernel_0080', + 'launch_dispatch_kernel_0081', + 'launch_dispatch_kernel_0082', + 'launch_dispatch_kernel_0083', + 'launch_dispatch_kernel_0084', + 'launch_dispatch_kernel_0085', + 'launch_dispatch_kernel_0086', + 'launch_dispatch_kernel_0087', + 'launch_dispatch_kernel_0088', + 'launch_dispatch_kernel_0089', + 'launch_dispatch_kernel_0090', + 'launch_dispatch_kernel_0091', + 'launch_dispatch_kernel_0092', + 'launch_dispatch_kernel_0093', + 'launch_dispatch_kernel_0094', + 'launch_dispatch_kernel_0095', + 'launch_dispatch_kernel_0096', + 'launch_dispatch_kernel_0097', + 'launch_dispatch_kernel_0098', + 'launch_dispatch_kernel_0099', + 'launch_dispatch_kernel_0100', + 'launch_dispatch_kernel_0101', + 'launch_dispatch_kernel_0102', + 'launch_dispatch_kernel_0103', + 'launch_dispatch_kernel_0104', + 'launch_dispatch_kernel_0105', + 'launch_dispatch_kernel_0106', + 'launch_dispatch_kernel_0107', + 'launch_dispatch_kernel_0108', + 'launch_dispatch_kernel_0109', + 'launch_dispatch_kernel_0110', + 'launch_dispatch_kernel_0111', + 'launch_dispatch_kernel_0112', + 'launch_dispatch_kernel_0113', + 'launch_dispatch_kernel_0114', + 'launch_dispatch_kernel_0115', + 'launch_dispatch_kernel_0116', + 'launch_dispatch_kernel_0117', + 'launch_dispatch_kernel_0118', + 'launch_dispatch_kernel_0119', + 'launch_dispatch_kernel_0120', + 'launch_dispatch_kernel_0121', + 'launch_dispatch_kernel_0122', + 'launch_dispatch_kernel_0123', + 'launch_dispatch_kernel_0124', + 'launch_dispatch_kernel_0125', + 'launch_dispatch_kernel_0126', + 'launch_dispatch_kernel_0127', + 'launch_dispatch_kernel_0128', + 'launch_dispatch_kernel_0129', + 'launch_dispatch_kernel_0130', + 'launch_dispatch_kernel_0131', + 'launch_dispatch_kernel_0132', + 'launch_dispatch_kernel_0133', + 'launch_dispatch_kernel_0134', + 'launch_dispatch_kernel_0135', + 'launch_dispatch_kernel_0136', + 'launch_dispatch_kernel_0137', + 'launch_dispatch_kernel_0138', + 'launch_dispatch_kernel_0139', + 'launch_dispatch_kernel_0140', + 'launch_dispatch_kernel_0141', + 'launch_dispatch_kernel_0142', + 'launch_dispatch_kernel_0143', + 'launch_dispatch_kernel_0144', + 'launch_dispatch_kernel_0145', + 'launch_dispatch_kernel_0146', + 'launch_dispatch_kernel_0147', + 'launch_dispatch_kernel_0148', + 'launch_dispatch_kernel_0149', + 'launch_dispatch_kernel_0150', + 'launch_dispatch_kernel_0151', + 'launch_dispatch_kernel_0152', + 'launch_dispatch_kernel_0153', + 'launch_dispatch_kernel_0154', + 'launch_dispatch_kernel_0155', + 'launch_dispatch_kernel_0156', + 'launch_dispatch_kernel_0157', + 'launch_dispatch_kernel_0158', + 'launch_dispatch_kernel_0159', + 'launch_dispatch_kernel_0160', + 'launch_dispatch_kernel_0161', + 'launch_dispatch_kernel_0162', + 'launch_dispatch_kernel_0163', + 'launch_dispatch_kernel_0164', + 'launch_dispatch_kernel_0165', + 'launch_dispatch_kernel_0166', + 'launch_dispatch_kernel_0167', + 'launch_dispatch_kernel_0168', + 'launch_dispatch_kernel_0169', + 'launch_dispatch_kernel_0170', + 'launch_dispatch_kernel_0171', + 'launch_dispatch_kernel_0172', + 'launch_dispatch_kernel_0173', + 'launch_dispatch_kernel_0174', + 'launch_dispatch_kernel_0175', + 'launch_dispatch_kernel_0176', + 'launch_dispatch_kernel_0177', + 'launch_dispatch_kernel_0178', + 'launch_dispatch_kernel_0179', + 'launch_dispatch_kernel_0180', + 'launch_dispatch_kernel_0181', + 'launch_dispatch_kernel_0182', + 'launch_dispatch_kernel_0183', + 'launch_dispatch_kernel_0184', + 'launch_dispatch_kernel_0185', + 'launch_dispatch_kernel_0186', + 'launch_dispatch_kernel_0187', + 'launch_dispatch_kernel_0188', + 'launch_dispatch_kernel_0189', + 'launch_dispatch_kernel_0190', + 'launch_dispatch_kernel_0191', + 'launch_dispatch_kernel_0192', + 'launch_dispatch_kernel_0193', + 'launch_dispatch_kernel_0194', + 'launch_dispatch_kernel_0195', + 'launch_dispatch_kernel_0196', + 'launch_dispatch_kernel_0197', + 'launch_dispatch_kernel_0198', + 'launch_dispatch_kernel_0199', + 'launch_dispatch_kernel_0200', + 'launch_dispatch_kernel_0201', + 'launch_dispatch_kernel_0202', + 'launch_dispatch_kernel_0203', + 'launch_dispatch_kernel_0204', + 'launch_dispatch_kernel_0205', + 'launch_dispatch_kernel_0206', + 'launch_dispatch_kernel_0207', + 'launch_dispatch_kernel_0208', + 'launch_dispatch_kernel_0209', + 'launch_dispatch_kernel_0210', + 'launch_dispatch_kernel_0211', + 'launch_dispatch_kernel_0212', + 'launch_dispatch_kernel_0213', + 'launch_dispatch_kernel_0214', + 'launch_dispatch_kernel_0215', + 'launch_dispatch_kernel_0216', + 'launch_dispatch_kernel_0217', + 'launch_dispatch_kernel_0218', + 'launch_dispatch_kernel_0219', + 'launch_dispatch_kernel_0220', + 'launch_dispatch_kernel_0221', + 'launch_dispatch_kernel_0222', + 'launch_dispatch_kernel_0223', + 'launch_dispatch_kernel_0224', + 'launch_dispatch_kernel_0225', + 'launch_dispatch_kernel_0226', + 'launch_dispatch_kernel_0227', + 'launch_dispatch_kernel_0228', + 'launch_dispatch_kernel_0229', + 'launch_dispatch_kernel_0230', + 'launch_dispatch_kernel_0231', + 'launch_dispatch_kernel_0232', + 'launch_dispatch_kernel_0233', + 'launch_dispatch_kernel_0234', + 'launch_dispatch_kernel_0235', + 'launch_dispatch_kernel_0236', + 'launch_dispatch_kernel_0237', + 'launch_dispatch_kernel_0238', + 'launch_dispatch_kernel_0239', + 'launch_dispatch_kernel_0240', + 'launch_dispatch_kernel_0241', + 'launch_dispatch_kernel_0242', + 'launch_dispatch_kernel_0243', + 'launch_dispatch_kernel_0244', + 'launch_dispatch_kernel_0245', + 'launch_dispatch_kernel_0246', + 'launch_dispatch_kernel_0247', + 'launch_dispatch_kernel_0248', + 'launch_dispatch_kernel_0249', + 'launch_dispatch_kernel_0250', + 'launch_dispatch_kernel_0251', + 'launch_dispatch_kernel_0252', + 'launch_dispatch_kernel_0253', + 'launch_dispatch_kernel_0254', + 'launch_dispatch_kernel_0255', + 'launch_dispatch_kernel_0256', + 'launch_dispatch_kernel_0257', + 'launch_dispatch_kernel_0258', + 'launch_dispatch_kernel_0259', + 'launch_dispatch_kernel_0260', + 'launch_dispatch_kernel_0261', + 'launch_dispatch_kernel_0262', + 'launch_dispatch_kernel_0263', + 'launch_dispatch_kernel_0264', + 'launch_dispatch_kernel_0265', + 'launch_dispatch_kernel_0266', + 'launch_dispatch_kernel_0267', + 'launch_dispatch_kernel_0268', + 'launch_dispatch_kernel_0269', + 'launch_dispatch_kernel_0270', + 'launch_dispatch_kernel_0271', + 'launch_dispatch_kernel_0272', + 'launch_dispatch_kernel_0273', + 'launch_dispatch_kernel_0274', + 'launch_dispatch_kernel_0275', + 'launch_dispatch_kernel_0276', + 'launch_dispatch_kernel_0277', + 'launch_dispatch_kernel_0278', + 'launch_dispatch_kernel_0279', + 'launch_dispatch_kernel_0280', + 'launch_dispatch_kernel_0281', + 'launch_dispatch_kernel_0282', + 'launch_dispatch_kernel_0283', + 'launch_dispatch_kernel_0284', + 'launch_dispatch_kernel_0285', + 'launch_dispatch_kernel_0286', + 'launch_dispatch_kernel_0287', + 'launch_dispatch_kernel_0288', + 'launch_dispatch_kernel_0289', + 'launch_dispatch_kernel_0290', + 'launch_dispatch_kernel_0291', + 'launch_dispatch_kernel_0292', + 'launch_dispatch_kernel_0293', + 'launch_dispatch_kernel_0294', + 'launch_dispatch_kernel_0295', + 'launch_dispatch_kernel_0296', + 'launch_dispatch_kernel_0297', + 'launch_dispatch_kernel_0298', + 'launch_dispatch_kernel_0299', + 'launch_dispatch_kernel_0300', + 'launch_dispatch_kernel_0301', + 'launch_dispatch_kernel_0302', + 'launch_dispatch_kernel_0303', + 'launch_dispatch_kernel_0304', + 'launch_dispatch_kernel_0305', + 'launch_dispatch_kernel_0306', + 'launch_dispatch_kernel_0307', + 'launch_dispatch_kernel_0308', + 'launch_dispatch_kernel_0309', + 'launch_dispatch_kernel_0310', + 'launch_dispatch_kernel_0311', + 'launch_dispatch_kernel_0312', + 'launch_dispatch_kernel_0313', + 'launch_dispatch_kernel_0314', + 'launch_dispatch_kernel_0315', + 'launch_dispatch_kernel_0316', + 'launch_dispatch_kernel_0317', + 'launch_dispatch_kernel_0318', + 'launch_dispatch_kernel_0319', + 'launch_dispatch_kernel_0320', + 'launch_dispatch_kernel_0321', + 'launch_dispatch_kernel_0322', + 'launch_dispatch_kernel_0323', + 'launch_dispatch_kernel_0324', + 'launch_dispatch_kernel_0325', + 'launch_dispatch_kernel_0326', + 'launch_dispatch_kernel_0327', + 'launch_dispatch_kernel_0328', + 'launch_dispatch_kernel_0329', + 'launch_dispatch_kernel_0330', + 'launch_dispatch_kernel_0331', + 'launch_dispatch_kernel_0332', + 'launch_dispatch_kernel_0333', + 'launch_dispatch_kernel_0334', + 'launch_dispatch_kernel_0335', + 'launch_dispatch_kernel_0336', + 'launch_dispatch_kernel_0337', + 'launch_dispatch_kernel_0338', + 'launch_dispatch_kernel_0339', + 'launch_dispatch_kernel_0340', + 'launch_dispatch_kernel_0341', + 'launch_dispatch_kernel_0342', + 'launch_dispatch_kernel_0343', + 'launch_dispatch_kernel_0344', + 'launch_dispatch_kernel_0345', + 'launch_dispatch_kernel_0346', + 'launch_dispatch_kernel_0347', + 'launch_dispatch_kernel_0348', + 'launch_dispatch_kernel_0349', + 'launch_dispatch_kernel_0350', + 'launch_dispatch_kernel_0351', + 'launch_dispatch_kernel_0352', + 'launch_dispatch_kernel_0353', + 'launch_dispatch_kernel_0354', + 'launch_dispatch_kernel_0355', + 'launch_dispatch_kernel_0356', + 'launch_dispatch_kernel_0357', + 'register_tvm_ffi', + 'tvm_ffi_function_names', +] + +# Semantic exports generated from export_plan.package_exports. +from .interface import FlashKMeansAssignRuntime as FlashKMeansAssignRuntime +from .interface import PreparedFlashKMeansAssign as PreparedFlashKMeansAssign +from .interface import init as init +from .interface import prepare_flash_kmeans_assign as prepare_flash_kmeans_assign +from .interface import flash_kmeans_assign_prepared as flash_kmeans_assign_prepared +from .interface import flash_kmeans_assign as flash_kmeans_assign +__all__ = [*globals().get('__all__', []), 'FlashKMeansAssignRuntime', 'PreparedFlashKMeansAssign', 'init', 'prepare_flash_kmeans_assign', 'flash_kmeans_assign_prepared', 'flash_kmeans_assign'] diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_benchmark.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_benchmark.py new file mode 100644 index 00000000..fd0e6677 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_benchmark.py @@ -0,0 +1,648 @@ +from __future__ import annotations + +import bisect +import contextlib +import ctypes +import importlib +import importlib.metadata +import math +import statistics +import sys +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +@dataclass(frozen=True) +class BenchResult: + """Strict CUPTI timing plus explicitly non-official host diagnostics. + + Cold-L2 flushing completes before host timestamps begin, so synchronized + E2E isolates semantic call start through candidate completion. ``times_ms`` + remains the official correlated GPU span. + """ + + times_ms: list[float] + backend: str = "cupti" + kernel_sum_times_ms: list[float] | None = None + inter_kernel_gap_times_ms: list[float] | None = None + active_union_times_ms: list[float] | None = None + activity_counts: list[int] | None = None + launch_activity_counts: list[int] | None = None + kernel_activity_counts: list[int] | None = None + submission_times_ms: list[float] | None = None + synchronized_e2e_times_ms: list[float] | None = None + cold_first_call_host_enqueue_ms: float | None = None + cold_first_call_synchronized_e2e_ms: float | None = None + + def __post_init__(self) -> None: + if self.backend != "cupti": + raise ValueError(f"exported benchmark timing backend must be 'cupti', got {self.backend!r}") + + @property + def median_ms(self) -> float: + return float(statistics.median(self.times_ms)) + + @property + def min_ms(self) -> float: + return float(min(self.times_ms)) + + @property + def mean_ms(self) -> float: + return float(statistics.fmean(self.times_ms)) + + @property + def median_gpu_span_ms(self) -> float: + return self.median_ms + + @property + def median_kernel_sum_ms(self) -> float | None: + if self.kernel_sum_times_ms is None: + return None + return float(statistics.median(self.kernel_sum_times_ms)) + + @property + def median_inter_kernel_gap_ms(self) -> float | None: + if self.inter_kernel_gap_times_ms is None: + return None + return float(statistics.median(self.inter_kernel_gap_times_ms)) + + @property + def median_active_union_ms(self) -> float | None: + if self.active_union_times_ms is None: + return None + return float(statistics.median(self.active_union_times_ms)) + + @property + def median_activity_count(self) -> float | None: + if self.activity_counts is None: + return None + return float(statistics.median(self.activity_counts)) + + @property + def median_launch_activity_count(self) -> float | None: + if self.launch_activity_counts is None: + return None + return float(statistics.median(self.launch_activity_counts)) + + @property + def median_kernel_activity_count(self) -> float | None: + if self.kernel_activity_counts is None: + return None + return float(statistics.median(self.kernel_activity_counts)) + + @property + def median_submission_ms(self) -> float | None: + if self.submission_times_ms is None: + return None + return float(statistics.median(self.submission_times_ms)) + + @property + def median_host_enqueue_ms(self) -> float | None: + return self.median_submission_ms + + @property + def host_enqueue_times_ms(self) -> list[float] | None: + return self.submission_times_ms + + @property + def median_synchronized_e2e_ms(self) -> float | None: + if self.synchronized_e2e_times_ms is None: + return None + return float(statistics.median(self.synchronized_e2e_times_ms)) + + +@dataclass(frozen=True) +class HostCallTiming: + """Diagnostic host brackets for one call; never an official GPU timing.""" + + host_enqueue_ms: float + synchronized_e2e_ms: float + + +@dataclass(frozen=True) +class _CuptiTiming: + gpu_span_ms: list[float] + kernel_sum_ms: list[float] + inter_kernel_gap_ms: list[float] + active_union_ms: list[float] + activity_count: list[int] + launch_activity_count: list[int] + kernel_activity_count: list[int] + + +class _L2Flusher: + def __init__(self) -> None: + import torch + + l2_size = int(torch.cuda.get_device_properties(0).L2_cache_size) + if l2_size <= 0: + raise RuntimeError("CUDA device did not report a positive L2 cache size") + self._buffer = torch.empty(2 * l2_size, dtype=torch.int8, device="cuda") + + def flush(self) -> None: + self._buffer.zero_() + + +_CUPTI: Any | None = None + + +def _extend_cuda_namespace_for_pathfinder() -> None: + try: + import cuda + except ImportError: + return + cuda_paths = getattr(cuda, "__path__", None) + if cuda_paths is None: + return + known = {str(path) for path in cuda_paths} + for entry in sys.path: + if not entry: + continue + cuda_root = Path(entry) / "cuda" + if (cuda_root / "pathfinder").is_dir() and str(cuda_root) not in known: + cuda_paths.append(str(cuda_root)) + known.add(str(cuda_root)) + + +def _preload_cupti_library() -> None: + # cupti-python imports cuda.pathfinder even when libcupti was found through + # the nvidia-cuda-cupti distribution. Extend the split CUDA namespace + # before either loading path so the later extension import is reliable. + _extend_cuda_namespace_for_pathfinder() + try: + distribution = importlib.metadata.distribution("nvidia-cuda-cupti") + major = importlib.metadata.version("cupti-python").split(".", 1)[0] + packaged = next( + ( + distribution.locate_file(path) + for path in distribution.files or () + if str(path).endswith(f"/libcupti.so.{major}") + ), + None, + ) + if packaged is not None and Path(packaged).is_file(): + ctypes.CDLL(str(packaged), mode=ctypes.RTLD_GLOBAL) + return + except (importlib.metadata.PackageNotFoundError, OSError): + pass + try: + pathfinder = importlib.import_module("cuda.pathfinder") + except ImportError: + return + loader = getattr(pathfinder, "load_nvidia_dynamic_lib", None) + if loader is None: + return + loaded = loader("cupti") + loaded_path = str(getattr(loaded, "abs_path", "")) + try: + major = importlib.metadata.version("cupti-python").split(".", 1)[0] + except importlib.metadata.PackageNotFoundError: + return + expected = f"libcupti.so.{major}" if major.isdigit() else None + if expected and loaded_path: + name = Path(loaded_path).name + if name.startswith("libcupti.so.") and name != expected: + raise ImportError(f"incompatible CUPTI library {loaded_path}; cupti-python expects {expected}") + + +def require_cupti() -> Any: + global _CUPTI + if _CUPTI is not None: + return _CUPTI + try: + _preload_cupti_library() + from cupti import cupti + except ImportError as exc: + raise RuntimeError( + "CUPTI timing is required; install the exported repository with its benchmark extra: " + "python -m pip install -e '.[benchmark]'" + ) from exc + _CUPTI = cupti + return _CUPTI + + +def measure_host_call(fn: Callable[[], Any]) -> tuple[Any, HostCallTiming]: + """Run one call and return explicitly labeled host diagnostics. + + The enqueue bracket ends when ``fn`` returns. The synchronized bracket + ends after ``torch.cuda.synchronize()``. A function that synchronizes + internally will therefore have a blocking enqueue bracket; callers should + not interpret either value as GPU-only execution time. + """ + import torch + + cupti = require_cupti() + start = cupti.get_timestamp() + value = fn() + submitted = cupti.get_timestamp() + torch.cuda.synchronize() + completed = cupti.get_timestamp() + return value, HostCallTiming( + host_enqueue_ms=(submitted - start) / 1e6, + synchronized_e2e_ms=(completed - start) / 1e6, + ) + + +def _finite_timing(value: Any, *, name: str, positive: bool) -> float: + number = float(value) + if not math.isfinite(number) or (number <= 0.0 if positive else number < 0.0): + relation = "positive" if positive else "non-negative" + raise ValueError(f"{name} must be finite and {relation}, got {value!r}") + return number + + +def _percentile(values: list[float], fraction: float) -> float: + if not values: + raise ValueError("percentile requires at least one value") + ordered = sorted(values) + rank = (len(ordered) - 1) * float(fraction) + lower = math.floor(rank) + upper = math.ceil(rank) + if lower == upper: + return ordered[lower] + weight = rank - lower + return ordered[lower] * (1.0 - weight) + ordered[upper] * weight + + +def _timing_distribution(values: Any, *, name: str, positive: bool) -> dict[str, Any]: + if not isinstance(values, list) or not values: + raise ValueError(f"{name} must contain at least one timing sample") + samples = [_finite_timing(value, name=name, positive=positive) for value in values] + return { + "sample_count": len(samples), + "min": min(samples), + "median": float(statistics.median(samples)), + "p90": _percentile(samples, 0.90), + "max": max(samples), + } + + +def _host_call_payload(timing: HostCallTiming, *, name: str) -> dict[str, Any]: + host_enqueue_ms = _finite_timing(timing.host_enqueue_ms, name=f"{name}.host_enqueue_ms", positive=False) + synchronized_e2e_ms = _finite_timing( + timing.synchronized_e2e_ms, + name=f"{name}.synchronized_e2e_ms", + positive=True, + ) + if host_enqueue_ms > synchronized_e2e_ms: + raise ValueError(f"{name}.host_enqueue_ms cannot exceed synchronized_e2e_ms") + return { + "timing_class": "cupti_timestamp_host_diagnostic", + "host_enqueue_ms": host_enqueue_ms, + "synchronized_e2e_ms": synchronized_e2e_ms, + } + + +def _hot_compute_payload(timing: BenchResult, *, name: str, cache_state: str) -> dict[str, Any]: + if timing.backend != "cupti": + raise ValueError(f"{name} must use CUPTI") + gpu_span = _timing_distribution(timing.times_ms, name=f"{name}.gpu_span_ms", positive=True) + host_enqueue = _timing_distribution( + timing.submission_times_ms, + name=f"{name}.host_enqueue_ms", + positive=False, + ) + synchronized_e2e = _timing_distribution( + timing.synchronized_e2e_times_ms, + name=f"{name}.synchronized_e2e_ms", + positive=True, + ) + sample_counts = { + gpu_span["sample_count"], + host_enqueue["sample_count"], + synchronized_e2e["sample_count"], + } + if len(sample_counts) != 1: + raise ValueError(f"{name} CUPTI/host timing sample counts must match") + for index, (enqueue_sample, e2e_sample) in enumerate( + zip(timing.submission_times_ms, timing.synchronized_e2e_times_ms, strict=True) + ): + if float(enqueue_sample) > float(e2e_sample): + raise ValueError(f"{name} host enqueue sample {index} cannot exceed its synchronized E2E sample") + return { + "timing_backend": "cupti", + "cache_state": cache_state, + "sample_count": gpu_span["sample_count"], + "gpu_span_ms": gpu_span, + "host_enqueue_ms": host_enqueue, + "synchronized_e2e_ms": synchronized_e2e, + } + + +def runtime_lifecycle_metrics( + *, + api: str, + measurement_session_id: str, + timing_boundary: str, + output_policy: str, + init: HostCallTiming | None, + init_sample_id: str | None, + first_compute: HostCallTiming, + first_cache_state: str, + hot_compute: BenchResult, + hot_cache_state: str, + amortization_call_counts: tuple[int, ...] = (1, 10, 100, 1000), + code_cache_state: str = "process_order_dependent", +) -> dict[str, Any]: + """Normalize init-once, first-signature, and repeated public-call evidence. + + ``first_compute`` is intentionally a CUPTI-timestamp host diagnostic: route + selection, compilation, allocation, launch, and completion all belong to + its synchronized E2E bracket, but it is not mislabeled as GPU-only time. + ``hot_compute`` retains strict correlated CUPTI activity timing. + """ + + for field_name, value in ( + ("api", api), + ("measurement_session_id", measurement_session_id), + ("timing_boundary", timing_boundary), + ("output_policy", output_policy), + ("first_cache_state", first_cache_state), + ("hot_cache_state", hot_cache_state), + ("code_cache_state", code_cache_state), + ): + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"runtime lifecycle {field_name} must be a non-empty string") + if ( + not amortization_call_counts + or any(isinstance(value, bool) or not isinstance(value, int) or value <= 0 for value in amortization_call_counts) + or tuple(amortization_call_counts) != tuple(sorted(set(amortization_call_counts))) + ): + raise ValueError("amortization_call_counts must be strictly increasing positive integers") + + init_payload = _host_call_payload(init, name="init_once") if init is not None else None + if init_payload is not None: + if not isinstance(init_sample_id, str) or not init_sample_id.strip(): + raise ValueError("init_sample_id is required when init timing is present") + init_payload = {"sample_id": init_sample_id, **init_payload} + elif init_sample_id is not None: + raise ValueError("init_sample_id requires init timing") + first_payload = { + "cache_state": first_cache_state, + "code_cache_state": code_cache_state, + **_host_call_payload(first_compute, name="first_compute"), + "gpu_span_ms": None, + "gpu_span_status": "not_collected_for_slot_miss_host_diagnostic", + } + hot_payload = _hot_compute_payload(hot_compute, name="hot_compute", cache_state=hot_cache_state) + + first_enqueue = first_payload["host_enqueue_ms"] + first_e2e = first_payload["synchronized_e2e_ms"] + hot_enqueue = hot_payload["host_enqueue_ms"]["median"] + hot_e2e = hot_payload["synchronized_e2e_ms"]["median"] + init_enqueue = init_payload["host_enqueue_ms"] if init_payload is not None else 0.0 + init_e2e = init_payload["synchronized_e2e_ms"] if init_payload is not None else 0.0 + amortized: list[dict[str, Any]] = [] + for call_count in amortization_call_counts: + repeated = call_count - 1 + after_init_enqueue = (first_enqueue + repeated * hot_enqueue) / call_count + after_init_e2e = (first_e2e + repeated * hot_e2e) / call_count + amortized.append( + { + "public_call_count": call_count, + "after_init_host_enqueue_ms_per_call": after_init_enqueue, + "after_init_synchronized_e2e_ms_per_call": after_init_e2e, + "including_init_host_enqueue_ms_per_call": ( + (init_enqueue + first_enqueue + repeated * hot_enqueue) / call_count + ), + "including_init_synchronized_e2e_ms_per_call": ( + (init_e2e + first_e2e + repeated * hot_e2e) / call_count + ), + } + ) + + return { + "schema": "loom-public-runtime-lifecycle-v1", + "api": api, + "measurement_session_id": measurement_session_id, + "timing_boundary": timing_boundary, + "output_policy": output_policy, + "init_once": init_payload, + "first_compute": first_payload, + "hot_compute": hot_payload, + "amortization": { + "model": "observed_first_call_plus_repeated_hot_median", + "init_attribution": "one_init_sample_per_validation_shard_runtime", + "missing_init_policy": "zero_for_api_without_explicit_init", + "call_counts": amortized, + }, + } + + +def compare_runtime_lifecycles(candidate: dict[str, Any], baseline: dict[str, Any]) -> dict[str, Any]: + """Compare two normalized public-API lifecycle records.""" + + for name, lifecycle in (("candidate", candidate), ("baseline", baseline)): + if not isinstance(lifecycle, dict) or lifecycle.get("schema") != "loom-public-runtime-lifecycle-v1": + raise ValueError(f"{name} runtime lifecycle has an invalid schema") + candidate_hot = candidate["hot_compute"] + baseline_hot = baseline["hot_compute"] + candidate_rows = candidate["amortization"]["call_counts"] + baseline_rows = baseline["amortization"]["call_counts"] + candidate_counts = [row["public_call_count"] for row in candidate_rows] + baseline_counts = [row["public_call_count"] for row in baseline_rows] + if candidate_counts != baseline_counts: + raise ValueError("candidate and baseline amortization call counts must match") + + amortized: list[dict[str, Any]] = [] + for candidate_row, baseline_row in zip(candidate_rows, baseline_rows, strict=True): + candidate_after_init = candidate_row["after_init_synchronized_e2e_ms_per_call"] + baseline_after_init = baseline_row["after_init_synchronized_e2e_ms_per_call"] + candidate_including_init = candidate_row["including_init_synchronized_e2e_ms_per_call"] + baseline_including_init = baseline_row["including_init_synchronized_e2e_ms_per_call"] + amortized.append( + { + "public_call_count": candidate_row["public_call_count"], + "after_init_synchronized_e2e_speedup": baseline_after_init / candidate_after_init, + "including_init_synchronized_e2e_speedup": baseline_including_init / candidate_including_init, + } + ) + + return { + "schema": "loom-public-runtime-lifecycle-comparison-v1", + "speedup_convention": "baseline_latency_divided_by_candidate_latency", + "hot_synchronized_e2e_speedup": ( + baseline_hot["synchronized_e2e_ms"]["median"] + / candidate_hot["synchronized_e2e_ms"]["median"] + ), + "hot_gpu_span_speedup": ( + baseline_hot["gpu_span_ms"]["median"] / candidate_hot["gpu_span_ms"]["median"] + ), + "amortized": amortized, + } + + +def _complete_l2_flush_before_bracket(flusher: Any, synchronize: Callable[[], None]) -> None: + """Finish cold-L2 preconditioning before host latency timestamps begin.""" + if flusher is None: + return + flusher.flush() + synchronize() + + +def _correlate( + cpu_brackets: list[tuple[int, int, int]], + launches: list[tuple[int, int, int]], + kernels: list[tuple[int, int, int]], +) -> _CuptiTiming: + if not launches or not kernels: + raise RuntimeError("CUPTI collected no launch/kernel activities") + launches.sort(key=lambda item: item[0]) + launch_starts = [item[0] for item in launches] + kernels.sort(key=lambda item: item[0]) + kernel_indices_by_correlation: dict[int, list[int]] = {} + for index, (_start, _end, correlation_id) in enumerate(kernels): + kernel_indices_by_correlation.setdefault(correlation_id, []).append(index) + + gpu_span_ms: list[float] = [] + kernel_sum_ms: list[float] = [] + inter_kernel_gap_ms: list[float] = [] + active_union_ms: list[float] = [] + activity_count: list[int] = [] + launch_activity_count: list[int] = [] + kernel_activity_count: list[int] = [] + for bracket_start, bracket_end, _completed in cpu_brackets: + lo = bisect.bisect_left(launch_starts, bracket_start) + hi = bisect.bisect_right(launch_starts, bracket_end) + correlation_ids = {launches[index][2] for index in range(lo, hi)} + launch_activity_count.append( + sum( + 1 + for index in range(lo, hi) + if kernel_indices_by_correlation.get(launches[index][2]) + ) + ) + selected_indices = { + index + for correlation_id in correlation_ids + for index in kernel_indices_by_correlation.get(correlation_id, ()) + } + if not selected_indices: + raise RuntimeError("CUPTI could not correlate a benchmark iteration with GPU kernel activity") + spans = [(kernels[index][0], kernels[index][1]) for index in selected_indices] + span_ns = max(end for _, end in spans) - min(start for start, _ in spans) + sum_ns = sum(end - start for start, end in spans) + covered_ns = 0 + current_start, current_end = sorted(spans)[0] + for start, end in sorted(spans)[1:]: + if start <= current_end: + current_end = max(current_end, end) + else: + covered_ns += current_end - current_start + current_start, current_end = start, end + covered_ns += current_end - current_start + gpu_span_ms.append(span_ns / 1e6) + kernel_sum_ms.append(sum_ns / 1e6) + active_union_ms.append(covered_ns / 1e6) + inter_kernel_gap_ms.append((span_ns - covered_ns) / 1e6) + activity_count.append(len(selected_indices)) + kernel_activity_count.append(len(selected_indices)) + return _CuptiTiming( + gpu_span_ms=gpu_span_ms, + kernel_sum_ms=kernel_sum_ms, + inter_kernel_gap_ms=inter_kernel_gap_ms, + active_union_ms=active_union_ms, + activity_count=activity_count, + launch_activity_count=launch_activity_count, + kernel_activity_count=kernel_activity_count, + ) + + +def bench_gpu_time( + fn: Callable[[], Any], + *, + warmup_iters: int = 5, + bench_iters: int = 20, + cold_l2: bool = True, + cold_first_call: HostCallTiming | None = None, +) -> BenchResult: + """Measure a zero-argument GPU workload with strict CUPTI activity tracing. + + L2 is flushed before every warmup and measured iteration. This function + never falls back to wall-clock or CUDA-event timing. + """ + if warmup_iters < 0 or bench_iters <= 0: + raise ValueError("warmup_iters must be non-negative and bench_iters must be positive") + + import torch + + cupti = require_cupti() + flusher = _L2Flusher() if cold_l2 else None + for _ in range(warmup_iters): + if flusher is not None: + flusher.flush() + fn() + torch.cuda.synchronize() + + launch_kinds = {int(cupti.ActivityKind.RUNTIME), int(cupti.ActivityKind.DRIVER)} + kernel_kinds = {int(cupti.ActivityKind.CONCURRENT_KERNEL)} + launches: list[tuple[int, int, int]] = [] + kernels: list[tuple[int, int, int]] = [] + + def _buffer_requested(): + return 8 * 1024 * 1024, 0 + + def _buffer_completed(activities: list[Any]) -> None: + for activity in activities: + record = (activity.start, activity.end, activity.correlation_id) + kind = int(activity.kind) + if kind in launch_kinds: + launches.append(record) + elif kind in kernel_kinds: + kernels.append(record) + + kinds = [ + cupti.ActivityKind.RUNTIME, + cupti.ActivityKind.DRIVER, + cupti.ActivityKind.CONCURRENT_KERNEL, + cupti.ActivityKind.MEMCPY, + cupti.ActivityKind.MEMSET, + ] + enabled: list[Any] = [] + cpu_brackets: list[tuple[int, int, int]] = [] + watchdog = sys.modules.get("loom.runtime.cuda_watchdog") + suspend_polling = getattr(watchdog, "suspend_polling", None) + polling_guard = suspend_polling() if callable(suspend_polling) else contextlib.nullcontext() + with polling_guard: + cupti.activity_register_callbacks(_buffer_requested, _buffer_completed) + try: + for kind in kinds: + cupti.activity_enable(kind) + enabled.append(kind) + + for _ in range(bench_iters): + _complete_l2_flush_before_bracket(flusher, torch.cuda.synchronize) + start = cupti.get_timestamp() + fn() + submitted = cupti.get_timestamp() + torch.cuda.synchronize() + completed = cupti.get_timestamp() + cpu_brackets.append((start, submitted, completed)) + cupti.activity_flush_all(1) + finally: + for kind in reversed(enabled): + cupti.activity_disable(kind) + cupti.finalize() + + timing = _correlate(cpu_brackets, launches, kernels) + return BenchResult( + times_ms=timing.gpu_span_ms, + kernel_sum_times_ms=timing.kernel_sum_ms, + inter_kernel_gap_times_ms=timing.inter_kernel_gap_ms, + active_union_times_ms=timing.active_union_ms, + activity_counts=timing.activity_count, + launch_activity_counts=timing.launch_activity_count, + kernel_activity_counts=timing.kernel_activity_count, + submission_times_ms=[(submitted - start) / 1e6 for start, submitted, _ in cpu_brackets], + synchronized_e2e_times_ms=[ + (completed - start) / 1e6 for start, _, completed in cpu_brackets + ], + cold_first_call_host_enqueue_ms=( + cold_first_call.host_enqueue_ms if cold_first_call is not None else None + ), + cold_first_call_synchronized_e2e_ms=( + cold_first_call.synchronized_e2e_ms if cold_first_call is not None else None + ), + ) diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/__init__.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_d160_padded_b23d_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_d160_padded_b23d_v1.py new file mode 100644 index 00000000..ca0579f3 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_d160_padded_b23d_v1.py @@ -0,0 +1,142 @@ +"""Clean-room Flash-KMeans Euclidean assignment D160 padded-tail candidate. + +Minimum architecture: sm_100a. This candidate uses a Weave BF16 pack/zero-fill +producer to pad logical D=160 tensors to D=192 scratch, then reuses the +Blackwell tcgen05/TMEM D192 three-chunk score/argmin path. It is not intended +for sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +from . import flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1 as _d192 +from .flash_kmeans_assign_stream import stream_cache_key +BLOCK_N = _d192.BLOCK_N +BLOCK_K = _d192.BLOCK_K +FEAT_D = 160 +FEAT_D_PAD = _d192.FEAT_D +FEAT_D_PAD_VECS = FEAT_D_PAD // 8 +PACK_THREADS = 256 +PACK_GRID_CAP = 4096 +_TMAP_CACHE: dict[tuple[int, int, int, int, int, int], Any] = {} +flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1", "arg_keys": ["x", "centroids", "x_pad", "c_pad", "B", "N", "D", "K", "total_x_pad", "total_c_pad"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 256}')) +pack_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1", "arg_keys": ["x", "centroids", "x_pad", "c_pad", "B", "N", "D", "K", "total_x_pad", "total_c_pad"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 256}')) + +def _cuda_include_dirs() -> list[str]: + return _d192._cuda_include_dirs() + +def _create_padded_tensor_map_3d(data_ptr: int, global_height: int, shared_height: int): + import torch + from .._dispatch_runtime import create_tensor_map_3d + device_index = torch.cuda.current_device() + key = (device_index, int(data_ptr), int(global_height), int(shared_height), FEAT_D_PAD, FEAT_D_PAD) + cached = _TMAP_CACHE.get(key) + if cached is not None: + return cached + cached = create_tensor_map_3d(data_ptr, global_height, shared_height, FEAT_D_PAD, FEAT_D_PAD).to(device=torch.device('cuda', device_index)) + _TMAP_CACHE[key] = cached + return cached + +def _make_tmaps(inputs: dict[str, Any], x_pad: Any, c_pad: Any) -> tuple[Any, Any]: + cache_key = (int(x_pad.data_ptr()), int(c_pad.data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), FEAT_D_PAD) + cached = inputs.get('_flash_kmeans_assign_cleanroom_d160_padded_b23d_v1_tmaps') + if cached is not None and cached[0] == cache_key: + return (cached[1], cached[2]) + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + tmap_x = _create_padded_tensor_map_3d(x_pad.data_ptr(), bsz * n_points, BLOCK_N) + tmap_c = _create_padded_tensor_map_3d(c_pad.data_ptr(), bsz * n_clusters, BLOCK_K) + inputs['_flash_kmeans_assign_cleanroom_d160_padded_b23d_v1_tmaps'] = (cache_key, tmap_x, tmap_c) + return (tmap_x, tmap_c) + +def _scratch_buffers(inputs: dict[str, Any]) -> tuple[Any, Any]: + import torch + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + cache_key = stream_cache_key(inputs, int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), bsz, n_points, n_clusters, int(inputs['D'])) + cache = inputs.setdefault('_flash_kmeans_assign_cleanroom_d160_padded_b23d_v1_scratch', {}) + cached = cache.get(cache_key) + if cached is not None: + return cached + x_pad = torch.empty((bsz, n_points, FEAT_D_PAD), dtype=inputs['x'].dtype, device=inputs['x'].device) + c_pad = torch.empty((bsz, n_clusters, FEAT_D_PAD), dtype=inputs['centroids'].dtype, device=inputs['centroids'].device) + cache[cache_key] = (x_pad, c_pad) + inputs.pop('_flash_kmeans_assign_cleanroom_d160_padded_b23d_v1_tmaps', None) + return (x_pad, c_pad) + +def _compiled_pack_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0357"}, "kernel_flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1", 0, 256]}')) + +def _launch_pack(inputs: dict[str, Any], x_pad: Any, c_pad: Any) -> None: + from .._dispatch_runtime import CUDAKernel + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + total_x_pad = bsz * n_points * FEAT_D_PAD + total_c_pad = bsz * n_clusters * FEAT_D_PAD + work_items = max(total_x_pad, total_c_pad) + grid_x = min((work_items + PACK_THREADS - 1) // PACK_THREADS, PACK_GRID_CAP) + cubin, kernel_name, smem_bytes, threads = _compiled_pack_kernel() + args = [inputs['x'], inputs['centroids'], x_pad, c_pad, bsz, n_points, dim, n_clusters, total_x_pad, total_c_pad] + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=(grid_x, 1, 1), block=(threads, 1, 1), args=args, shared_mem=smem_bytes) + +def _launch_d192_main_on_padded(inputs: dict[str, Any], tmap_x: Any, tmap_c: Any) -> None: + from .._dispatch_runtime import CUDAKernel + from .._dispatch_runtime import pack_kernel_args + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + cubin, kernel_name, smem_bytes, threads = _d192._compiled_kernel() + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // BLOCK_K + grid = (bsz * num_n_tiles, 1, 1) + block = (threads, 1, 1) + args = pack_kernel_args(_d192.flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1, x_tmap=tmap_x, c_tmap=tmap_c, x_sq=inputs['x_sq'], c_sq=inputs['c_sq'], out=inputs['out'], B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles) + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=grid, block=block, args=args, shared_mem=smem_bytes) + +def launch_for_eval(inputs: dict[str, Any]) -> Any: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + dtype = str(inputs.get('dtype', getattr(inputs['x'], 'dtype', 'bfloat16'))).replace('torch.', '') + if dtype not in {'bfloat16', 'bf16'}: + raise ValueError(''.join(['flash_kmeans_assign_cleanroom_tcgen05_d160_padded_b23d_v1 requires bfloat16, got ', format(dtype, '')])) + if dim != FEAT_D: + raise ValueError(''.join(['flash_kmeans_assign_cleanroom_tcgen05_d160_padded_b23d_v1 requires D=', format(FEAT_D, ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(n_clusters, '')])) + x_pad, c_pad = _scratch_buffers(inputs) + _launch_pack(inputs, x_pad, c_pad) + tmap_x, tmap_c = _make_tmaps(inputs, x_pad, c_pad) + _launch_d192_main_on_padded(inputs, tmap_x, tmap_c) + return {'cluster_ids': inputs['out']} + +def compile_and_launch_flash_kmeans_assign_cleanroom(B: int=1, N: int=2048, K: int=2048, D: int=160, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(16001) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 16001}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1.py new file mode 100644 index 00000000..8fa679f0 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1.py @@ -0,0 +1,91 @@ +"""Clean-room Flash-KMeans Euclidean assignment D192 repeated-MMA single-tile candidate. + +Minimum architecture: sm_100a. This candidate uses Blackwell tcgen05/TMEM +(``lm.mma``) for the D-specific repeated-MMA score tile producer and is not intended for +sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +BLOCK_N = 128 +BLOCK_K = 256 +SCORE_CHUNK_K = 128 +CSQ_VEC = 4 +CSQ_STAGE_VEC = 4 +FEAT_D = 192 +NUM_COMPUTE_WARPS = 4 +X_TILE_BYTES = BLOCK_N * FEAT_D * 2 +C_TILE_BYTES = BLOCK_K * FEAT_D * 2 +X_FEAT_CHUNK_BYTES = BLOCK_N * 64 * 2 +C_FEAT_CHUNK_BYTES = BLOCK_K * 64 * 2 +CSQ_TILE_BYTES = BLOCK_K * 4 +flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 149504, "constants": [], "cta_group": 1, "threads": 192}')) + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as _common_cuda_include_dirs + return _common_cuda_include_dirs() + +def _make_tmaps(inputs: dict[str, Any]) -> tuple[Any, Any]: + cache_key = (int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), int(inputs['D'])) + cached = inputs.get('_flash_kmeans_assign_cleanroom_d192_single_v1_tmaps') + if cached is not None and cached[0] == cache_key: + return (cached[1], cached[2]) + from .._dispatch_runtime import create_tensor_map_3d + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + dim = int(inputs['D']) + tmap_x = create_tensor_map_3d(inputs['x'].data_ptr(), bsz * n_points, BLOCK_N, dim, dim) + tmap_c = create_tensor_map_3d(inputs['centroids'].data_ptr(), bsz * n_clusters, BLOCK_K, dim, dim) + inputs['_flash_kmeans_assign_cleanroom_d192_single_v1_tmaps'] = (cache_key, tmap_x, tmap_c) + return (tmap_x, tmap_c) + +def _compiled_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0344"}, "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1", 149504, 192]}')) + +def launch_for_eval(inputs: dict[str, Any]) -> Any: + from .._dispatch_runtime import CUDAKernel + from .._dispatch_runtime import pack_kernel_args + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + if dim != FEAT_D: + raise ValueError(''.join(['flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1 requires D=', format(FEAT_D, ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(n_clusters, '')])) + tmap_x, tmap_c = _make_tmaps(inputs) + cubin, kernel_name, smem_bytes, threads = _compiled_kernel() + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // BLOCK_K + grid = (bsz * num_n_tiles, 1, 1) + block = (threads, 1, 1) + args = pack_kernel_args(flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1, x_tmap=tmap_x, c_tmap=tmap_c, x_sq=inputs['x_sq'], c_sq=inputs['c_sq'], out=inputs['out'], B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles) + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=grid, block=block, args=args, shared_mem=smem_bytes) + return {'cluster_ids': inputs['out']} + +def compile_and_launch_flash_kmeans_assign_cleanroom(B: int=2, N: int=1024, K: int=512, D: int=192, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(5101) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'B': B, 'N': N, 'D': D, 'K': K, 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + passed = bool(torch.equal(out, ref)) + result: dict[str, Any] = {'passed': passed} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 5101}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1.py new file mode 100644 index 00000000..6fadc933 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1.py @@ -0,0 +1,106 @@ +"""Clean-room Flash-KMeans Euclidean assignment D192 repeated-MMA paired-tile candidate. + +Minimum architecture: sm_100a. This candidate dispatches between a D-specific repeated-MMA single point-tile path and a D-specific repeated-MMA paired point-tile tcgen05/TMEM path. It is not intended for +sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +from . import flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1 as _single +BLOCK_N = 128 +BLOCK_K = 256 +SCORE_CHUNK_K = 128 +CSQ_VEC = 16 +CSQ_STAGE_VEC = 4 +FEAT_D = 192 +NUM_COMPUTE_WARPS = 4 +X_TILE_BYTES = BLOCK_N * FEAT_D * 2 +C_TILE_BYTES = BLOCK_K * FEAT_D * 2 +X_FEAT_CHUNK_BYTES = BLOCK_N * 64 * 2 +C_FEAT_CHUNK_BYTES = BLOCK_K * 64 * 2 +CSQ_TILE_BYTES = BLOCK_K * 4 +X1_OFFSET = X_TILE_BYTES +C_OFFSET = 2 * X_TILE_BYTES +CSQ_OFFSET = 2 * X_TILE_BYTES + C_TILE_BYTES +PAIRED_GRID_CAP = 2516 +flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 198656, "constants": [], "cta_group": 1, "threads": 192}')) + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as _common_cuda_include_dirs + return _common_cuda_include_dirs() + +def _make_tmaps(inputs: dict[str, Any]) -> tuple[Any, Any]: + cache_key = (int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), int(inputs['D'])) + cached = inputs.get('_flash_kmeans_assign_cleanroom_d192_splitd_v1_tmaps') + if cached is not None and cached[0] == cache_key: + return (cached[1], cached[2]) + from .._dispatch_runtime import create_tensor_map_3d + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + dim = int(inputs['D']) + tmap_x = create_tensor_map_3d(inputs['x'].data_ptr(), bsz * n_points, BLOCK_N, dim, dim) + tmap_c = create_tensor_map_3d(inputs['centroids'].data_ptr(), bsz * n_clusters, BLOCK_K, dim, dim) + inputs['_flash_kmeans_assign_cleanroom_d192_splitd_v1_tmaps'] = (cache_key, tmap_x, tmap_c) + return (tmap_x, tmap_c) + +def _compiled_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0343"}, "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", 198656, 192]}')) +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 198656, "constants": [], "cta_group": 1, "threads": 192}')) + +def _use_single_tile_path(*, n_points: int, n_clusters: int) -> bool: + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // BLOCK_K + return num_n_tiles % 2 != 0 or (num_n_tiles <= 8 and k_tiles <= 2) + +def launch_for_eval(inputs: dict[str, Any]) -> Any: + from .._dispatch_runtime import CUDAKernel + from .._dispatch_runtime import pack_kernel_args + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + if _use_single_tile_path(n_points=n_points, n_clusters=n_clusters): + return _single.launch_for_eval(inputs) + if dim != FEAT_D: + raise ValueError(''.join(['flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1 requires D=', format(FEAT_D, ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(n_clusters, '')])) + num_n_tiles = n_points // BLOCK_N + if num_n_tiles % 2 != 0: + raise ValueError('flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1 requires an even number of BLOCK_N point tiles') + tmap_x, tmap_c = _make_tmaps(inputs) + cubin, kernel_name, smem_bytes, threads = _compiled_kernel() + k_tiles = n_clusters // BLOCK_K + total_tiles = bsz * (num_n_tiles // 2) + grid = (min(total_tiles, PAIRED_GRID_CAP), 1, 1) + block = (threads, 1, 1) + args = pack_kernel_args(ir, x_tmap=tmap_x, c_tmap=tmap_c, x_sq=inputs['x_sq'], c_sq=inputs['c_sq'], out=inputs['out'], B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles) + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=grid, block=block, args=args, shared_mem=smem_bytes) + return {'cluster_ids': inputs['out']} + +def compile_and_launch_flash_kmeans_assign_cleanroom(B: int=2, N: int=1024, K: int=512, D: int=192, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(5101) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'B': B, 'N': N, 'D': D, 'K': K, 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + passed = bool(torch.equal(out, ref)) + result: dict[str, Any] = {'passed': passed} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 5101}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1.py new file mode 100644 index 00000000..68271344 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1.py @@ -0,0 +1,91 @@ +"""Clean-room Flash-KMeans Euclidean assignment D256 repeated-MMA single-tile candidate. + +Minimum architecture: sm_100a. This candidate uses Blackwell tcgen05/TMEM +(``lm.mma``) for the D-specific repeated-MMA score tile producer and is not intended for +sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +BLOCK_N = 128 +BLOCK_K = 256 +SCORE_CHUNK_K = 128 +CSQ_VEC = 4 +CSQ_STAGE_VEC = 4 +FEAT_D = 256 +NUM_COMPUTE_WARPS = 4 +X_TILE_BYTES = BLOCK_N * FEAT_D * 2 +C_TILE_BYTES = BLOCK_K * FEAT_D * 2 +X_FEAT_CHUNK_BYTES = BLOCK_N * 64 * 2 +C_FEAT_CHUNK_BYTES = BLOCK_K * 64 * 2 +CSQ_TILE_BYTES = BLOCK_K * 4 +flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 198656, "constants": [], "cta_group": 1, "threads": 192}')) + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as _common_cuda_include_dirs + return _common_cuda_include_dirs() + +def _make_tmaps(inputs: dict[str, Any]) -> tuple[Any, Any]: + cache_key = (int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), int(inputs['D'])) + cached = inputs.get('_flash_kmeans_assign_cleanroom_d256_single_v1_tmaps') + if cached is not None and cached[0] == cache_key: + return (cached[1], cached[2]) + from .._dispatch_runtime import create_tensor_map_3d + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + dim = int(inputs['D']) + tmap_x = create_tensor_map_3d(inputs['x'].data_ptr(), bsz * n_points, BLOCK_N, dim, dim) + tmap_c = create_tensor_map_3d(inputs['centroids'].data_ptr(), bsz * n_clusters, BLOCK_K, dim, dim) + inputs['_flash_kmeans_assign_cleanroom_d256_single_v1_tmaps'] = (cache_key, tmap_x, tmap_c) + return (tmap_x, tmap_c) + +def _compiled_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0342"}, "kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", 198656, 192]}')) + +def launch_for_eval(inputs: dict[str, Any]) -> Any: + from .._dispatch_runtime import CUDAKernel + from .._dispatch_runtime import pack_kernel_args + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + if dim != FEAT_D: + raise ValueError(''.join(['flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1 requires D=', format(FEAT_D, ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(n_clusters, '')])) + tmap_x, tmap_c = _make_tmaps(inputs) + cubin, kernel_name, smem_bytes, threads = _compiled_kernel() + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // BLOCK_K + grid = (bsz * num_n_tiles, 1, 1) + block = (threads, 1, 1) + args = pack_kernel_args(flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1, x_tmap=tmap_x, c_tmap=tmap_c, x_sq=inputs['x_sq'], c_sq=inputs['c_sq'], out=inputs['out'], B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles) + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=grid, block=block, args=args, shared_mem=smem_bytes) + return {'cluster_ids': inputs['out']} + +def compile_and_launch_flash_kmeans_assign_cleanroom(B: int=2, N: int=1024, K: int=512, D: int=256, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(5101) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'B': B, 'N': N, 'D': D, 'K': K, 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + passed = bool(torch.equal(out, ref)) + result: dict[str, Any] = {'passed': passed} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 5101}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_b23d_v2.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_b23d_v2.py new file mode 100644 index 00000000..d02c4154 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_b23d_v2.py @@ -0,0 +1,90 @@ +"""Flash-KMeans Euclidean assignment D160/D192/D256 repeated-MMA bucket seed. + +Minimum architecture: sm_100a. This wrapper routes contract evaluation to +D-specific Weave tcgen05/TMEM seeds only; it does not call external runtime +fallbacks. The seed modules are not intended for sm_120a/sm_121a where ptxas +rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from typing import Any +from . import flash_kmeans_assign_cleanroom_tcgen05_d160_padded_b23d_v1 as _d160 +from . import flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1 as _d192 +from . import flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1 as _d256 +BLOCK_N = _d192.BLOCK_N +BLOCK_K = _d192.BLOCK_K +SUPPORTED_DIMS = (160, 192, 256) +D160_BUCKET_SHAPES: list[dict[str, Any]] = [{'label': 'd160_paired_b1_n2048_k2048_d160', 'params': {'B': 1, 'N': 2048, 'D': 160, 'K': 2048, 'dtype': 'bfloat16', 'seed': 16001}}, {'label': 'd160_fallback_b2_n2432_k1024_d160', 'params': {'B': 2, 'N': 2432, 'D': 160, 'K': 1024, 'dtype': 'bfloat16', 'seed': 16002}}] +BUCKET_SHAPES: list[dict[str, Any]] = [*D160_BUCKET_SHAPES, {'label': 'd192_small_b1_n1024_k512_d192', 'params': {'B': 1, 'N': 1024, 'D': 192, 'K': 512, 'dtype': 'bfloat16', 'seed': 19201}}, {'label': 'd192_paired_b1_n2048_k2048_d192', 'params': {'B': 1, 'N': 2048, 'D': 192, 'K': 2048, 'dtype': 'bfloat16', 'seed': 19202}}, {'label': 'd192_random_b2_n4096_k768_d192', 'params': {'B': 2, 'N': 4096, 'D': 192, 'K': 768, 'dtype': 'bfloat16', 'seed': 19204}}, {'label': 'd256_small_b1_n1024_k512_d256', 'params': {'B': 1, 'N': 1024, 'D': 256, 'K': 512, 'dtype': 'bfloat16', 'seed': 25601}}, {'label': 'd256_paired_b1_n2048_k2048_d256', 'params': {'B': 1, 'N': 2048, 'D': 256, 'K': 2048, 'dtype': 'bfloat16', 'seed': 25602}}, {'label': 'd256_fallback_b2_n2432_k1024_d256', 'params': {'B': 2, 'N': 2432, 'D': 256, 'K': 1024, 'dtype': 'bfloat16', 'seed': 25603}}, {'label': 'd256_large_b4_n16384_k1024_d256', 'params': {'B': 4, 'N': 16384, 'D': 256, 'K': 1024, 'dtype': 'bfloat16', 'seed': 25604}}, {'label': 'd256_hugek_b1_n1024_k8192_d256', 'params': {'B': 1, 'N': 1024, 'D': 256, 'K': 8192, 'dtype': 'bfloat16', 'seed': 25605}}] + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + dim = int(inputs['D']) + module = _module_for_dim(dim) + route_kind = _route_kind(inputs) + outputs = module.launch_for_eval(inputs) + normalized = _normalize_outputs(outputs, inputs) + selected_route = ''.join(['d', format(dim, ''), '_', format(route_kind, ''), '_repeated_mma_v2']) + normalized['selected_route'] = selected_route + normalized['route_trace'] = {'shape_key': _shape_key(inputs), 'selected_route': selected_route, 'selected_entrypoint': ''.join([format(module.__name__, ''), ':launch_for_eval']), 'selected_seed': ''.join(['d', format(dim, ''), '-', format(route_kind, ''), '-repeated-mma-b23d-v2']), 'route_kind': 'specialized', 'route_source': 'shape-specific-seed', 'guard_id': ''.join(['guard_d', format(dim, ''), '_', format(route_kind, ''), '_repeated_mma_b23d_v2']), 'guard_condition': ''.join(['dtype == bfloat16 and D == ', format(dim, ''), ' and N % ', format(BLOCK_N, ''), ' == 0 and K % ', format(BLOCK_K, ''), ' == 0']), 'classification': 'seed-probe', 'dispatcher_kernel_ms': None, 'reason': 'non-128D bucket seed wrapper routes to D-specific Weave tcgen05 candidate'} + return normalized + +def _module_for_dim(dim: int): + if dim == 160: + return _d160 + if dim == 192: + return _d192 + if dim == 256: + return _d256 + raise ValueError(''.join(['non128d split-D seed supports D in ', format(SUPPORTED_DIMS, ''), ', got ', format(dim, '')])) + +def _route_kind(inputs: dict[str, Any]) -> str: + dim = int(inputs['D']) + if dim == 160: + return 'padded_single' + if dim == 256: + return 'single' + n_tiles = int(inputs['N']) // BLOCK_N + k_tiles = int(inputs['K']) // BLOCK_K + if n_tiles % 2 != 0 or (n_tiles <= 8 and k_tiles <= 2): + return 'single' + return 'paired' + +def _normalize_outputs(outputs: Any, inputs: dict[str, Any]) -> dict[str, Any]: + if outputs is None: + return {'cluster_ids': inputs['out']} + if hasattr(outputs, 'shape'): + return {'cluster_ids': outputs} + if isinstance(outputs, dict): + normalized = dict(outputs) + if 'cluster_ids' not in normalized and 'out' in normalized: + normalized['cluster_ids'] = normalized['out'] + if 'cluster_ids' in normalized: + return normalized + raise TypeError("flash_kmeans_assign non128d split-D route must return cluster_ids or write inputs['out']") + +def _shape_key(inputs: dict[str, Any]) -> str: + label = inputs.get('label') + if label: + return str(label) + return ''.join(['b', format(int(inputs['B']), ''), '_n', format(int(inputs['N']), ''), '_k', format(int(inputs['K']), ''), '_d', format(int(inputs['D']), '')]) + +def compile_and_launch_flash_kmeans_assign_cleanroom(B: int=1, N: int=2048, K: int=2048, D: int=160, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(16001) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 16001}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_v1.py new file mode 100644 index 00000000..5ed326dc --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_v1.py @@ -0,0 +1,96 @@ +"""Flash-KMeans Euclidean assignment D192/D256 repeated-MMA bucket seed. + +Minimum architecture: sm_100a. This wrapper routes contract evaluation to +D-specific Weave tcgen05/TMEM seeds only; it does not call external runtime +fallbacks. The seed modules are not intended for sm_120a/sm_121a where ptxas +rejects tcgen05 instructions. D160 is listed as a blocked sibling because its +32-wide BF16 tail currently hangs when lowered as K=32 tcgen05 MMA. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from typing import Any +from . import flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1 as _d192_single +from . import flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1 as _d192 +from . import flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1 as _d256 +BLOCK_N = _d192.BLOCK_N +BLOCK_K = _d192.BLOCK_K +SUPPORTED_DIMS = (192, 256) +BUCKET_SHAPES: list[dict[str, Any]] = [{'label': 'd192_small_b2_n1024_k512_d192', 'params': {'B': 2, 'N': 1024, 'D': 192, 'K': 512, 'dtype': 'bfloat16', 'seed': 19201}}, {'label': 'd192_paired_b2_n2048_k2048_d192', 'params': {'B': 2, 'N': 2048, 'D': 192, 'K': 2048, 'dtype': 'bfloat16', 'seed': 19202}}, {'label': 'd192_fallback_b1_n2176_k1024_d192', 'params': {'B': 1, 'N': 2176, 'D': 192, 'K': 1024, 'dtype': 'bfloat16', 'seed': 19203}}, {'label': 'd192_large_b4_n32768_k1024_d192', 'params': {'B': 4, 'N': 32768, 'D': 192, 'K': 1024, 'dtype': 'bfloat16', 'seed': 19204}}, {'label': 'd256_small_b1_n1024_k512_d256', 'params': {'B': 1, 'N': 1024, 'D': 256, 'K': 512, 'dtype': 'bfloat16', 'seed': 25601}}, {'label': 'd256_paired_b1_n4096_k4096_d256', 'params': {'B': 1, 'N': 4096, 'D': 256, 'K': 4096, 'dtype': 'bfloat16', 'seed': 25602}}, {'label': 'd256_fallback_b2_n2432_k2048_d256', 'params': {'B': 2, 'N': 2432, 'D': 256, 'K': 2048, 'dtype': 'bfloat16', 'seed': 25603}}, {'label': 'd256_hugek_b1_n512_k8192_d256', 'params': {'B': 1, 'N': 512, 'D': 256, 'K': 8192, 'dtype': 'bfloat16', 'seed': 25604}}] +D160_BLOCKED_SHAPES: list[dict[str, Any]] = [{'label': 'd160_fallback_b2_n2432_k1024_d160', 'params': {'B': 2, 'N': 2432, 'D': 160, 'K': 1024, 'dtype': 'bfloat16', 'seed': 16002}, 'blocker': 'D160 requires a BF16 32-wide feature tail; the current K=32 tcgen05 path hangs at runtime.'}, {'label': 'd160_paired_b1_n2048_k2048_d160', 'params': {'B': 1, 'N': 2048, 'D': 160, 'K': 2048, 'dtype': 'bfloat16', 'seed': 16001}, 'blocker': 'D160 requires a BF16 32-wide feature tail; the current K=32 tcgen05 path hangs at runtime.'}] + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + dim = int(inputs['D']) + module = _module_for_dim(dim) + route_kind = _route_kind(inputs) + outputs = module.launch_for_eval(inputs) + normalized = _normalize_outputs(outputs, inputs) + selected_route = ''.join(['d', format(dim, ''), '_', format(route_kind, ''), '_repeated_mma_v1']) + selected_entrypoint = _selected_entrypoint(dim, route_kind) + normalized['selected_route'] = selected_route + normalized['route_trace'] = {'shape_key': _shape_key(inputs), 'selected_route': selected_route, 'selected_entrypoint': selected_entrypoint, 'selected_seed': ''.join(['d', format(dim, ''), '-', format(route_kind, ''), '-repeated-mma-v1']), 'route_kind': 'specialized', 'route_source': 'shape-specific-seed', 'guard_id': ''.join(['guard_d', format(dim, ''), '_', format(route_kind, ''), '_repeated_mma_v1']), 'guard_condition': ''.join(['dtype == bfloat16 and D == ', format(dim, ''), ' and N % ', format(BLOCK_N, ''), ' == 0 and K % ', format(BLOCK_K, ''), ' == 0']), 'classification': 'seed-probe', 'dispatcher_kernel_ms': None, 'reason': 'non-128D bucket seed wrapper routes to D-specific Weave tcgen05 candidate'} + return normalized + +def _module_for_dim(dim: int): + if dim == 192: + return _d192 + if dim == 256: + return _d256 + raise ValueError(''.join(['non128d split-D seed supports D in ', format(SUPPORTED_DIMS, ''), ', got ', format(dim, '')])) + +def _selected_entrypoint(dim: int, route_kind: str) -> str: + if dim == 192 and route_kind == 'single': + return ''.join([format(_d192_single.__name__, ''), ':launch_for_eval']) + if dim == 192: + return ''.join([format(_d192.__name__, ''), ':launch_for_eval']) + if dim == 256: + return ''.join([format(_d256.__name__, ''), ':launch_for_eval']) + raise ValueError(''.join(['non128d split-D seed supports D in ', format(SUPPORTED_DIMS, ''), ', got ', format(dim, '')])) + +def _route_kind(inputs: dict[str, Any]) -> str: + if int(inputs['D']) == 256: + return 'single' + n_tiles = int(inputs['N']) // BLOCK_N + k_tiles = int(inputs['K']) // BLOCK_K + if n_tiles % 2 != 0 or (n_tiles <= 8 and k_tiles <= 2): + return 'single' + return 'paired' + +def _normalize_outputs(outputs: Any, inputs: dict[str, Any]) -> dict[str, Any]: + if outputs is None: + return {'cluster_ids': inputs['out']} + if hasattr(outputs, 'shape'): + return {'cluster_ids': outputs} + if isinstance(outputs, dict): + normalized = dict(outputs) + if 'cluster_ids' not in normalized and 'out' in normalized: + normalized['cluster_ids'] = normalized['out'] + if 'cluster_ids' in normalized: + return normalized + raise TypeError("flash_kmeans_assign non128d split-D route must return cluster_ids or write inputs['out']") + +def _shape_key(inputs: dict[str, Any]) -> str: + label = inputs.get('label') + if label: + return str(label) + return ''.join(['b', format(int(inputs['B']), ''), '_n', format(int(inputs['N']), ''), '_k', format(int(inputs['K']), ''), '_d', format(int(inputs['D']), '')]) + +def compile_and_launch_flash_kmeans_assign_cleanroom(B: int=1, N: int=1024, K: int=512, D: int=192, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(5101) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 5101}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_v10.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_v10.py new file mode 100644 index 00000000..3483fd2c --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_v10.py @@ -0,0 +1,89 @@ +"""Clean-room Flash-KMeans Euclidean assignment candidate. + +Minimum architecture: sm_100a. This candidate uses Blackwell tcgen05/TMEM +(``lm.mma``) for the D=128 score tile producer and is not intended for +sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +BLOCK_N = 128 +BLOCK_K = 256 +SCORE_CHUNK_K = 128 +CSQ_VEC = 4 +CSQ_STAGE_VEC = 4 +FEAT_D = 128 +NUM_COMPUTE_WARPS = 4 +X_TILE_BYTES = BLOCK_N * FEAT_D * 2 +C_TILE_BYTES = BLOCK_K * FEAT_D * 2 +CSQ_TILE_BYTES = BLOCK_K * 4 +flash_kmeans_assign_cleanroom_tcgen05_v10 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_cleanroom_tcgen05_v10", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 100352, "constants": [], "cta_group": 1, "threads": 192}')) + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as _common_cuda_include_dirs + return _common_cuda_include_dirs() + +def _make_tmaps(inputs: dict[str, Any]) -> tuple[Any, Any]: + cache_key = (int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), int(inputs['D'])) + cached = inputs.get('_flash_kmeans_assign_cleanroom_v10_tmaps') + if cached is not None and cached[0] == cache_key: + return (cached[1], cached[2]) + from .._dispatch_runtime import create_tensor_map_3d + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + dim = int(inputs['D']) + tmap_x = create_tensor_map_3d(inputs['x'].data_ptr(), bsz * n_points, BLOCK_N, dim, dim) + tmap_c = create_tensor_map_3d(inputs['centroids'].data_ptr(), bsz * n_clusters, BLOCK_K, dim, dim) + inputs['_flash_kmeans_assign_cleanroom_v10_tmaps'] = (cache_key, tmap_x, tmap_c) + return (tmap_x, tmap_c) + +def _compiled_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0341"}, "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", 100352, 192]}')) + +def launch_for_eval(inputs: dict[str, Any]) -> Any: + from .._dispatch_runtime import CUDAKernel + from .._dispatch_runtime import pack_kernel_args + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + if dim != FEAT_D: + raise ValueError(''.join(['flash_kmeans_assign_cleanroom_tcgen05_v10 requires D=', format(FEAT_D, ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(n_clusters, '')])) + tmap_x, tmap_c = _make_tmaps(inputs) + cubin, kernel_name, smem_bytes, threads = _compiled_kernel() + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // BLOCK_K + grid = (bsz * num_n_tiles, 1, 1) + block = (threads, 1, 1) + args = pack_kernel_args(flash_kmeans_assign_cleanroom_tcgen05_v10, x_tmap=tmap_x, c_tmap=tmap_c, x_sq=inputs['x_sq'], c_sq=inputs['c_sq'], out=inputs['out'], B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles) + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=grid, block=block, args=args, shared_mem=smem_bytes) + return {'cluster_ids': inputs['out']} + +def compile_and_launch_flash_kmeans_assign_cleanroom(B: int=2, N: int=1024, K: int=512, D: int=128, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(5101) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'B': B, 'N': N, 'D': D, 'K': K, 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + passed = bool(torch.equal(out, ref)) + result: dict[str, Any] = {'passed': passed} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 5101}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_v15.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_v15.py new file mode 100644 index 00000000..23284712 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_cleanroom_tcgen05_v15.py @@ -0,0 +1,106 @@ +"""Clean-room Flash-KMeans Euclidean assignment 16-wide argmin candidate. + +Minimum architecture: sm_100a. This candidate dispatches between the v10 +single point-tile path for small grids and a 16-wide argmin variant of the +paired point-tile tcgen05/TMEM path for large grids. It is not intended for +sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +from . import flash_kmeans_assign_cleanroom_tcgen05_v10 as _single +BLOCK_N = 128 +BLOCK_K = 256 +SCORE_CHUNK_K = 128 +CSQ_VEC = 16 +CSQ_STAGE_VEC = 4 +FEAT_D = 128 +NUM_COMPUTE_WARPS = 4 +X_TILE_BYTES = BLOCK_N * FEAT_D * 2 +C_TILE_BYTES = BLOCK_K * FEAT_D * 2 +CSQ_TILE_BYTES = BLOCK_K * 4 +X1_OFFSET = X_TILE_BYTES +C_OFFSET = 2 * X_TILE_BYTES +CSQ_OFFSET = 2 * X_TILE_BYTES + C_TILE_BYTES +PAIRED_GRID_CAP = 2516 +flash_kmeans_assign_cleanroom_tcgen05_v15 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_cleanroom_tcgen05_v15", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 133120, "constants": [], "cta_group": 1, "threads": 192}')) + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as _common_cuda_include_dirs + return _common_cuda_include_dirs() + +def _make_tmaps(inputs: dict[str, Any]) -> tuple[Any, Any]: + cache_key = (int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), int(inputs['D'])) + cached = inputs.get('_flash_kmeans_assign_cleanroom_v15_tmaps') + if cached is not None and cached[0] == cache_key: + return (cached[1], cached[2]) + from .._dispatch_runtime import create_tensor_map_3d + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + dim = int(inputs['D']) + tmap_x = create_tensor_map_3d(inputs['x'].data_ptr(), bsz * n_points, BLOCK_N, dim, dim) + tmap_c = create_tensor_map_3d(inputs['centroids'].data_ptr(), bsz * n_clusters, BLOCK_K, dim, dim) + inputs['_flash_kmeans_assign_cleanroom_v15_tmaps'] = (cache_key, tmap_x, tmap_c) + return (tmap_x, tmap_c) + +def _compiled_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0340"}, "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", 133120, 192]}')) +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_cleanroom_tcgen05_v15", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 133120, "constants": [], "cta_group": 1, "threads": 192}')) + +def _use_single_tile_path(*, n_points: int, n_clusters: int) -> bool: + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // BLOCK_K + return num_n_tiles <= 8 and k_tiles <= 2 + +def launch_for_eval(inputs: dict[str, Any]) -> Any: + from .._dispatch_runtime import CUDAKernel + from .._dispatch_runtime import pack_kernel_args + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + if _use_single_tile_path(n_points=n_points, n_clusters=n_clusters): + return _single.launch_for_eval(inputs) + if dim != FEAT_D: + raise ValueError(''.join(['flash_kmeans_assign_cleanroom_tcgen05_v15 requires D=', format(FEAT_D, ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(n_clusters, '')])) + num_n_tiles = n_points // BLOCK_N + if num_n_tiles % 2 != 0: + raise ValueError('flash_kmeans_assign_cleanroom_tcgen05_v15 requires an even number of BLOCK_N point tiles') + tmap_x, tmap_c = _make_tmaps(inputs) + cubin, kernel_name, smem_bytes, threads = _compiled_kernel() + k_tiles = n_clusters // BLOCK_K + total_tiles = bsz * (num_n_tiles // 2) + grid = (min(total_tiles, PAIRED_GRID_CAP), 1, 1) + block = (threads, 1, 1) + args = pack_kernel_args(ir, x_tmap=tmap_x, c_tmap=tmap_c, x_sq=inputs['x_sq'], c_sq=inputs['c_sq'], out=inputs['out'], B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles) + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=grid, block=block, args=args, shared_mem=smem_bytes) + return {'cluster_ids': inputs['out']} + +def compile_and_launch_flash_kmeans_assign_cleanroom(B: int=2, N: int=1024, K: int=512, D: int=128, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(5101) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'B': B, 'N': N, 'D': D, 'K': K, 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + passed = bool(torch.equal(out, ref)) + result: dict[str, Any] = {'passed': passed} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 5101}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_4549_c72b_k1024_overlay_weave_evolve_flash_kmeans_assign_3ae6_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_4549_c72b_k1024_overlay_weave_evolve_flash_kmeans_assign_3ae6_v1.py new file mode 100644 index 00000000..912bfee9 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_4549_c72b_k1024_overlay_weave_evolve_flash_kmeans_assign_3ae6_v1.py @@ -0,0 +1,60 @@ +"""Frozen D112 K-owned portfolio composition (minimum architecture: sm_100a). + +The c72b dual-issuer tcgen05/TMEM kernel owns exactly K=1024. The 4549 +direct-handoff portfolio owns K in {256, 512, 768, 4096, 8192}. Both child +sources are imported byte-for-byte from their pinned source revisions; this +module adds only the host ownership guard and auditable route metadata. Every +route writes caller-owned ``cluster_ids`` with one Weave compute launch and no +padding, packing, fallback, or global workspace. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from typing import Any +from . import flash_kmeans_assign_d112_direct_handoff_full_bucket_weave_evolve_flash_kmeans_assign_4549_v1 as _base +from . import flash_kmeans_assign_d112_shared_point_dual_issuer_tcgen05_weave_evolve_flash_kmeans_assign_6e1e_v1 as _overlay +FEAT_D = 112 +BASE_K = (256, 512, 768, 4096, 8192) +OVERLAY_K = 1024 +SUPPORTED_K = (*BASE_K[:3], OVERLAY_K, *BASE_K[3:]) +ROUTE_ID = 'd112_4549_c72b_k1024_overlay_weave_evolve_flash_kmeans_assign_3ae6_v1' +SEED_ID = 'd112-4549-c72b-k1024-overlay-weave-evolve-flash-kmeans-assign-3ae6-v1' +TARGET_SHAPE = 'physical_D112_padded_14' +BASE_SOURCE = {'id': 'base_portfolio', 'source_task': 'weave-evolve-flash-kmeans-assign-4549', 'source_kernel_task': 'weave-evolve-flash-kmeans-assign-4549', 'source_revision': '21c48b897873f6a5e631c7fccd320c35d73ae6d5', 'evidence_revision': '713ee91855069de8580152c20ec37b5ab4569c29', 'source_blob': '9d4da373afd9869c2718fffaf62f4230f4bb2482', 'source_sha256': 'ef1b1b574a5ce0516c745a5c28fa381a47e8a5e5c9ad49e35a93bff496c43b64', 'source_path': 'loom/examples/weave/flash_kmeans_assign_d112_direct_handoff_full_bucket_weave_evolve_flash_kmeans_assign_4549_v1.py', 'entrypoint': ''.join([format(_base.__name__, ''), ':launch_for_eval'])} +OVERLAY_SOURCE = {'id': 'k1024_overlay', 'source_task': 'weave-evolve-flash-kmeans-assign-c72b', 'source_kernel_task': 'weave-evolve-flash-kmeans-assign-6e1e', 'source_revision': 'a2415111e8f5d268760c97ccc65496173b289741', 'evidence_revision': '048037f7fada3aec532d102d1f1ca16fd1decc8b', 'source_blob': 'cf307b65bed35cd7384fbb78de2bfcf10bc29193', 'source_sha256': '2cf0d0da5e68abe23c22cf5a88315f8aede639095e6b84df40931a423aa94766', 'source_path': 'loom/examples/weave/flash_kmeans_assign_d112_shared_point_dual_issuer_tcgen05_weave_evolve_flash_kmeans_assign_6e1e_v1.py', 'entrypoint': ''.join([format(_overlay.__name__, ''), ':launch_for_eval'])} + +def source_for_k(K: int) -> dict[str, str]: + """Return the immutable source record for one admitted K regime.""" + if K == OVERLAY_K: + return OVERLAY_SOURCE + if K in BASE_K: + return BASE_SOURCE + raise ValueError('D112 4549/c72b portfolio requires K in (256, 512, 768, 1024, 4096, 8192)') + +def supports_shape(*, D: int, N: int, K: int) -> bool: + if D != FEAT_D: + return False + if K == OVERLAY_K: + return _overlay.supports_shape(D=D, N=N, K=K) + if K in BASE_K: + return _base.supports_shape(D=D, N=N, K=K) + return False + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + B, N, D, K = (int(inputs[key]) for key in ('B', 'N', 'D', 'K')) + if not supports_shape(D=D, N=N, K=K): + raise ValueError('D112 4549/c72b portfolio requires D=112, an owned N multiple, and K in (256, 512, 768, 1024, 4096, 8192)') + source = source_for_k(K) + child = _overlay if K == OVERLAY_K else _base + caller_out = inputs['out'] + caller_ptr = int(caller_out.data_ptr()) + child_outputs = child.launch_for_eval(inputs) + cluster_ids = child_outputs.get('cluster_ids', child_outputs.get('out')) + if cluster_ids is None: + raise ValueError(''.join([format(source['id'], ''), ' did not return cluster_ids or out'])) + returned_ptr = int(cluster_ids.data_ptr()) + if returned_ptr != caller_ptr: + raise ValueError(''.join([format(source['id'], ''), " did not return caller-owned inputs['out']"])) + child_trace = dict(child_outputs.get('route_trace') or {}) + wrapper_entrypoint = ''.join([format(__name__, ''), ':launch_for_eval']) + return {'cluster_ids': cluster_ids, 'selected_route': ROUTE_ID, 'route_trace': {'shape_key': TARGET_SHAPE, 'selected_route': ROUTE_ID, 'selected_entrypoint': wrapper_entrypoint, 'selected_seed': SEED_ID, 'selected_source': source['id'], 'selected_source_task': source['source_task'], 'selected_source_kernel_task': source['source_kernel_task'], 'selected_source_revision': source['source_revision'], 'selected_source_evidence_revision': source['evidence_revision'], 'selected_source_blob': source['source_blob'], 'selected_source_sha256': source['source_sha256'], 'selected_source_path': source['source_path'], 'selected_source_entrypoint': source['entrypoint'], 'child_selected_route': child_outputs.get('selected_route'), 'child_route_trace': child_trace, 'dispatch_key': 'K', 'dispatch_value': K, 'owned_k': [OVERLAY_K] if K == OVERLAY_K else list(BASE_K), 'caller_output_data_ptr': caller_ptr, 'returned_output_data_ptr': returned_ptr, 'caller_owned_output': True, 'output_dtype': 'int32', 'compute_kernel_count': 1, 'padding_count': 0, 'pack_count': 0, 'fallback_contract_regions': [], 'residual_contract_regions': [], 'global_workspace': False, 'materialized_padding': False, 'production_route': 'weave_only', 'launch_grid': child_trace.get('launch_grid'), 'producer_to_consumer': child_trace.get('producer_to_consumer'), 'B': B, 'N': N, 'D': D, 'K': K}} diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_weave_evolve_flash_kmeans_assign_e5e1_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_weave_evolve_flash_kmeans_assign_e5e1_v1.py new file mode 100644 index 00000000..9c4b70f4 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_weave_evolve_flash_kmeans_assign_e5e1_v1.py @@ -0,0 +1,61 @@ +"""D112 K1024 direct consumer-scoped MMA handoff (minimum architecture: sm_90a). + +Eight CTA ranks own disjoint K/8 spans and two four-warp consumer groups own +disjoint halves of each span. TMA writes each group's two canonical 8x112 +centroid stages directly. The owning consumer group loads those stages with +``ldmatrix.x2`` into the ``mma.sync`` B fragment, so there is no +canonical-to-MMA-view materialization. Independent 128-thread named barriers +protect stage reuse without a 256-thread per-tile rendezvous. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +BLOCK_N, BLOCK_K, FEAT_D, THREADS = (64, 8, 112, 256) +CLUSTER_CTAS, CONSUMER_GROUPS, PIPELINE_STAGES = (8, 2, 2) +SUPPORTED_K = (1024,) +X_ELEMS, C_ELEMS = (BLOCK_N * FEAT_D, BLOCK_K * FEAT_D) +X_BYTES, C_BYTES = (X_ELEMS * 2, C_ELEMS * 2) +CENTROID_STAGE_COUNT = CONSUMER_GROUPS * PIPELINE_STAGES +SCORE_BYTES = CONSUMER_GROUPS * 4 * 16 * BLOCK_K * 4 +LOCAL_KEY_BYTES = BLOCK_N * 8 +GROUP_KEY_BYTES = CONSUMER_GROUPS * LOCAL_KEY_BYTES +CLUSTER_KEY_BYTES = CLUSTER_CTAS * LOCAL_KEY_BYTES +U32_MASK = 4294967295 +X_RAW_OFFSET = 0 +C_DIRECT_OFFSET = X_RAW_OFFSET + X_BYTES +X_MMA_OFFSET = C_DIRECT_OFFSET + CENTROID_STAGE_COUNT * C_BYTES +SCORE_OFFSET = X_MMA_OFFSET + X_BYTES +GROUP_KEY_OFFSET = SCORE_OFFSET + SCORE_BYTES +LOCAL_KEY_OFFSET = GROUP_KEY_OFFSET + GROUP_KEY_BYTES +CLUSTER_KEY_OFFSET = LOCAL_KEY_OFFSET + LOCAL_KEY_BYTES +SMEM_BYTES = CLUSTER_KEY_OFFSET + CLUSTER_KEY_BYTES +ROUTE_ID = 'd112_consumer_scoped_direct_mma_view_handoff_weave_evolve_flash_kmeans_assign_e5e1_v1' +SEED_ID = 'd112-consumer-scoped-direct-mma-view-handoff-weave-evolve-flash-kmeans-assign-e5e1-v1' +TARGET_SHAPE = 'physical_D112_padded_14' +flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1", "arg_keys": ["x", "centroids", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles"], "cluster_dims": [8, 1, 1], "computed_smem_bytes": 46592, "constants": [], "cta_group": 1, "threads": 256}')) +cluster_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1", "arg_keys": ["x", "centroids", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles"], "cluster_dims": [8, 1, 1], "computed_smem_bytes": 46592, "constants": [], "cta_group": 1, "threads": 256}')) +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1", "arg_keys": ["x", "centroids", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles"], "cluster_dims": [8, 1, 1], "computed_smem_bytes": 46592, "constants": [], "cta_group": 1, "threads": 256}')) + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as dirs + return dirs() + +def _compiled_cluster() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0356"}, "kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1", 46592, 256]}')) + +def supports_shape(*, D: int, N: int, K: int) -> bool: + return D == FEAT_D and N % BLOCK_N == 0 and (K in SUPPORTED_K) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + import torch + from .._dispatch_runtime import CUDAKernel + B, N, D, K = (int(inputs[key]) for key in ('B', 'N', 'D', 'K')) + if not supports_shape(D=D, N=N, K=K): + raise ValueError('consumer-scoped direct MMA handoff e5e1 requires D=112, N%64=0, and K=1024') + cubin, name, smem, threads = _compiled_cluster() + grid = (B * (N // BLOCK_N) * CLUSTER_CTAS, 1, 1) + with CUDAKernel(cubin, name) as kernel: + kernel.launch_cluster(grid=grid, block=(threads, 1, 1), args=[inputs['x'], inputs['centroids'], inputs['c_sq'], inputs['out'], B, N, D, K, N // BLOCK_N], cluster_dims=(CLUSTER_CTAS, 1, 1), shared_mem=smem, stream=torch.cuda.current_stream(), timeout_ms=120000) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': {'shape_key': TARGET_SHAPE, 'selected_route': ROUTE_ID, 'selected_entrypoint': ''.join([format(__name__, ''), ':launch_for_eval']), 'selected_seed': SEED_ID, 'route_kind': 'eight_owner_consumer_scoped_direct_mma_view_handoff', 'cluster_ctas': CLUSTER_CTAS, 'consumer_groups_per_owner': CONSUMER_GROUPS, 'centroid_pipeline_stages': PIPELINE_STAGES, 'centroids_per_owner': K // CLUSTER_CTAS, 'centroids_per_consumer_group': K // (CLUSTER_CTAS * CONSUMER_GROUPS), 'producer_ownership': 'ranks 0..7 stage points locally; each four-warp group owns two canonical centroid stages', 'reuse_mechanism': 'direct TMA canonical stages consumed by ldmatrix.x2 with independent 128-thread group barriers', 'producer_to_consumer': 'TMA_to_consumer_scoped_canonical_stage_to_ldmatrix_x2_to_mma_sync_to_argmax_to_onchip_cluster_ids', 'overlap_edge': 'centroid_t_plus_1_direct_TMA -> current_centroid_t_mma_and_argmax', 'stage_storage_bytes': C_BYTES, 'split_k_reduction': 'shared-memory intra-owner key merge followed by on-chip cluster merge; no global workspace', 'launch_grid': grid, 'residual_contract_regions': []}} diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_direct_handoff_full_bucket_weave_evolve_flash_kmeans_assign_4549_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_direct_handoff_full_bucket_weave_evolve_flash_kmeans_assign_4549_v1.py new file mode 100644 index 00000000..ced7cfee --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_direct_handoff_full_bucket_weave_evolve_flash_kmeans_assign_4549_v1.py @@ -0,0 +1,43 @@ +"""D112 K-general direct consumer handoff (minimum architecture: sm_90a). + +This additive full-bucket entrypoint retains the admitted e5e1 Weave IR: +direct canonical centroid stages feed ``ldmatrix.x2`` and ``mma.sync`` before +the owner-local argmax/key merge writes caller-owned cluster IDs. The lane is +compiled and validated explicitly for ``sm_100a``. Every admitted K is a +multiple of the eight CTA owners, two consumer groups, and the K8 MMA tile, so +the same dynamic loop topology covers the complete 14-row D112 bucket. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from typing import Any +from . import flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_weave_evolve_flash_kmeans_assign_e5e1_v1 as _direct +BLOCK_N = _direct.BLOCK_N +BLOCK_K = _direct.BLOCK_K +FEAT_D = _direct.FEAT_D +THREADS = _direct.THREADS +CLUSTER_CTAS = _direct.CLUSTER_CTAS +CONSUMER_GROUPS = _direct.CONSUMER_GROUPS +PIPELINE_STAGES = _direct.PIPELINE_STAGES +SUPPORTED_K = (256, 512, 768, 1024, 4096, 8192) +K_PARTITION_QUANTUM = CLUSTER_CTAS * CONSUMER_GROUPS * BLOCK_K +ROUTE_ID = 'd112_direct_handoff_full_bucket_weave_evolve_flash_kmeans_assign_4549_v1' +SEED_ID = 'd112-direct-handoff-full-bucket-weave-evolve-flash-kmeans-assign-4549-v1' +TARGET_SHAPE = 'physical_D112_padded_14' +cluster_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1", "arg_keys": ["x", "centroids", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles"], "cluster_dims": [8, 1, 1], "computed_smem_bytes": 46592, "constants": [], "cta_group": 1, "threads": 256}')) +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1", "arg_keys": ["x", "centroids", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles"], "cluster_dims": [8, 1, 1], "computed_smem_bytes": 46592, "constants": [], "cta_group": 1, "threads": 256}')) + +def supports_shape(*, D: int, N: int, K: int) -> bool: + return D == FEAT_D and N % BLOCK_N == 0 and (K in SUPPORTED_K) and (K >= K_PARTITION_QUANTUM) and (K % K_PARTITION_QUANTUM == 0) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + import torch + from .._dispatch_runtime import CUDAKernel + B, N, D, K = (int(inputs[key]) for key in ('B', 'N', 'D', 'K')) + if not supports_shape(D=D, N=N, K=K): + raise ValueError('D112 direct-handoff full bucket requires D=112, N%64=0, and K in (256, 512, 768, 1024, 4096, 8192)') + cubin, name, smem, threads = _direct._compiled_cluster() + grid = (B * (N // BLOCK_N) * CLUSTER_CTAS, 1, 1) + with CUDAKernel(cubin, name) as kernel: + kernel.launch_cluster(grid=grid, block=(threads, 1, 1), args=[inputs['x'], inputs['centroids'], inputs['c_sq'], inputs['out'], B, N, D, K, N // BLOCK_N], cluster_dims=(CLUSTER_CTAS, 1, 1), shared_mem=smem, stream=torch.cuda.current_stream(), timeout_ms=120000) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': {'shape_key': TARGET_SHAPE, 'selected_route': ROUTE_ID, 'selected_entrypoint': ''.join([format(__name__, ''), ':launch_for_eval']), 'selected_seed': SEED_ID, 'route_kind': 'eight_owner_k_general_consumer_scoped_direct_mma_view_handoff', 'cluster_ctas': CLUSTER_CTAS, 'consumer_groups_per_owner': CONSUMER_GROUPS, 'centroid_pipeline_stages': PIPELINE_STAGES, 'centroids_per_owner': K // CLUSTER_CTAS, 'centroids_per_consumer_group': K // (CLUSTER_CTAS * CONSUMER_GROUPS), 'producer_ownership': 'ranks 0..7 stage points locally; each four-warp group owns two canonical centroid stages', 'reuse_mechanism': 'direct TMA canonical stages consumed by ldmatrix.x2 with independent 128-thread group barriers', 'producer_to_consumer': 'TMA_to_consumer_scoped_canonical_stage_to_ldmatrix_x2_to_mma_sync_to_argmax_to_onchip_cluster_ids', 'overlap_edge': 'centroid_t_plus_1_direct_TMA -> current_centroid_t_mma_and_argmax', 'stage_storage_bytes': _direct.C_BYTES, 'split_k_reduction': 'shared-memory intra-owner key merge followed by on-chip cluster merge; no global workspace', 'launch_grid': grid, 'supported_k': SUPPORTED_K, 'residual_contract_regions': []}} diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1.py new file mode 100644 index 00000000..278e2044 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1.py @@ -0,0 +1,154 @@ +"""Exact-guard D112 c829/f826/4549 portfolio (minimum: sm_100a). + +The two assigned B4/D112/K1024 rows select the frozen c829 owner-local +warp-MMA kernel. B5/N2944/D112/K512 retains the frozen f826 two-owner +warp-MMA kernel through the 32b2 portfolio, and every other row in the exact +14-row D112 bucket retains the frozen 4549 child through that same portfolio. +The host-only composition adds no padding, packing, workspace, fallback, or +extra GPU launch. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from dataclasses import dataclass +from threading import RLock +from typing import Any, Callable +from . import flash_kmeans_assign_d112_f826_exact_portfolio_generalize_auto_tuning_flash_kmeans_assign_8d82_v1 as _retained +from . import flash_kmeans_assign_d112_k1024_owner_local_warp_mma_distinct_workfeed_weave_evolve_flash_kmeans_assign_c829_v1 as _c829 +FEAT_D = 112 +ROUTE_ID = 'd112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1' +SEED_ID = 'd112-f826-c829-full-bucket-weave-evolve-flash-kmeans-assign-a262-v1' +TARGET_SHAPE = 'physical_D112_padded_14' +SHAPE_KEYS = {(2, 2048, 112, 512): 'post_d895_d112_b2_n2048_k512_d112', (4, 8192, 112, 1024): 'post_d895_d112_b4_n8192_k1024_d112', (1, 512, 112, 8192): 'post_d895_d112_b1_n512_k8192_d112', (1, 256, 112, 256): 'post_d895_d112_b1_n256_k256_d112', (3, 3840, 112, 768): 'adjacent_9ca0_d112_tail_div_b3_n3840_k768_d112', (2, 512, 112, 8192): 'adjacent_3328_d112_highk_low_n_b2_n512_k8192_d112', (5, 2176, 112, 512): 'adjacent_5600_d112_odd_tile_tail_b5_n2176_k512_d112', (4, 3456, 112, 256): 'adjacent_a2f8_d112_lowk_tail_b4_n3456_k256_d112', (1, 384, 112, 4096): 'adjacent_8f09_d112_small_highk_b1_n384_k4096_d112', (3, 768, 112, 8192): 'adjacent_d9d5_d112_highk_request_b3_n768_k8192_d112', (5, 2944, 112, 512): 'adjacent_68cf_d112_tail_b5_n2944_k512_d112', (2, 3200, 112, 768): 'adjacent_c44f_d112_odd_tail_b2_n3200_k768_d112', (1, 1408, 112, 4096): 'adjacent_1d49_d112_highk_tail_b1_n1408_k4096_d112', (4, 3712, 112, 1024): 'adjacent_ef00_d112_odd_tail_b4_n3712_k1024_d112'} +EXACT_C829_KEYS = frozenset({(4, 8192, 112, 1024), (4, 3712, 112, 1024)}) +EXACT_F826_KEY = (5, 2944, 112, 512) +C829_SOURCE = {'id': 'c829_exact_k1024', 'source_task': 'weave-evolve-flash-kmeans-assign-ab9a', 'implementation_revision': 'e3d609c6c0e68796f1c9273a6e113205c5fcea1e', 'representative_measurement_revision': '2def42e0b313de745896952a8f17d1f3c78680a8', 'timed_source_blob': '6bb66ffdeb5da35094c3e2ee02932ee1cac2432e', 'timed_source_sha256': '3205f28d937bf9ff1616bbb0eb35e8765f13447a9802cb9807dc28d46b1ab27f', 'executable_body_blob': '9d0b7cb69e1274850fc60be2f6f1f78002f2b23f', 'executable_body_sha256': '72cc676d4ddcb32c40cabb01e7f2ff9ffc1138d8449ee49b38ec552c7a54b720', 'final_source_blob': '4ce799909b9b93a64c272c831b55ad4c13186fb3', 'final_source_sha256': '81bf56e0b8cbb8962f9852159096fa1df52c769350fd46412e782c78a7ded010', 'source_path': 'loom/examples/weave/flash_kmeans_assign_d112_k1024_owner_local_warp_mma_distinct_workfeed_weave_evolve_flash_kmeans_assign_c829_v1.py', 'entrypoint': ''.join([format(_c829.__name__, ''), ':launch_for_eval'])} +F826_SOURCE = _decode_capture(_json_loads('{"__dict_items__": [["id", "f826_exact_k512"], ["source_task", "weave-evolve-flash-kmeans-assign-f826"], ["implementation_commit", "124ef2ad24c3ac86c063e69da0cb0c5d7fca0174"], ["implementation_source_blob", "3a1dee841855c15481a740f28ba07142e4fdbb29"], ["implementation_source_sha256", "99c44085f258a699b7e3372115cc82aa969643addae0797bf47c0758cad8cf0f"], ["measurement_source_blob", "a3dcf64c5014d223127bf0a39b13c485cf140c48"], ["measurement_source_sha256", "fd3e97727367d3acbbd3b994a86cb25f563ce3f32544a724dbd09f05a70a88ce"], ["final_source_blob", "f30b546e70bfbc80a94547458d857ead0de5e320"], ["final_source_sha256", "47c2d22bbb00b3061d268fae9c405e62257bde8ba43ba33fc1124d456572cd63"], ["source_path", "loom/examples/weave/flash_kmeans_assign_d112_k512_two_owner_peer_only_mma_weave_evolve_flash_kmeans_assign_f826_v1.py"], ["entrypoint", "loom.examples.weave.flash_kmeans_assign_d112_k512_two_owner_peer_only_mma_weave_evolve_flash_kmeans_assign_f826_v1:launch_for_eval"]]}')) +BASE_4549_SOURCE = _decode_capture(_json_loads('{"__dict_items__": [["id", "base_portfolio"], ["source_task", "weave-evolve-flash-kmeans-assign-4549"], ["source_kernel_task", "weave-evolve-flash-kmeans-assign-4549"], ["source_revision", "21c48b897873f6a5e631c7fccd320c35d73ae6d5"], ["evidence_revision", "713ee91855069de8580152c20ec37b5ab4569c29"], ["source_blob", "9d4da373afd9869c2718fffaf62f4230f4bb2482"], ["source_sha256", "ef1b1b574a5ce0516c745a5c28fa381a47e8a5e5c9ad49e35a93bff496c43b64"], ["source_path", "loom/examples/weave/flash_kmeans_assign_d112_direct_handoff_full_bucket_weave_evolve_flash_kmeans_assign_4549_v1.py"], ["entrypoint", "loom.examples.weave.flash_kmeans_assign_d112_direct_handoff_full_bucket_weave_evolve_flash_kmeans_assign_4549_v1:launch_for_eval"]]}')) +RETAINED_PORTFOLIO_SOURCE = {'id': 'portfolio_32b2', 'source_task': 'generalize-auto-tuning-flash-kmeans-assign-32b2', 'source_blob': '2acba116338583b0f55b06b49b1882779b440b09', 'source_sha256': 'ba9da56a4ff50c211d201dcbfbb69d8d1cb31464c2628dd5110be1320b626087', 'source_path': 'loom/examples/weave/flash_kmeans_assign_d112_f826_exact_portfolio_generalize_auto_tuning_flash_kmeans_assign_8d82_v1.py', 'entrypoint': ''.join([format(_retained.__name__, ''), ':launch_for_eval'])} +_PREPARE_LOCK = RLock() + +def _shape(inputs: dict[str, Any]) -> tuple[int, int, int, int]: + return tuple((int(inputs[name]) for name in ('B', 'N', 'D', 'K'))) + +def _shape_key(key: tuple[int, int, int, int]) -> str: + try: + return SHAPE_KEYS[key] + except KeyError as exc: + raise ValueError(''.join(['a262 portfolio rejects unassigned shape ', format(repr(key), '')])) from exc + +def supports_shape(*, B: int, D: int, N: int, K: int) -> bool: + return (int(B), int(N), int(D), int(K)) in SHAPE_KEYS + +def source_for_shape(*, B: int, N: int, D: int, K: int) -> dict[str, Any]: + key = (int(B), int(N), int(D), int(K)) + _shape_key(key) + if key in EXACT_C829_KEYS: + return dict(C829_SOURCE) + return dict(_retained.source_for_shape(B=B, N=N, D=D, K=K)) + +def _resolve_route(inputs: dict[str, Any], *, force_c829: bool=False, force_f826: bool=False) -> tuple[Callable[[dict[str, Any]], dict[str, Any]], dict[str, Any], str, str]: + key = _shape(inputs) + _shape_key(key) + if force_c829 and force_f826: + raise ValueError('only one forced route may be requested') + if force_c829 and key not in EXACT_C829_KEYS: + raise ValueError('forced c829 requires one of the two exact K1024 keys') + if force_f826 and key != EXACT_F826_KEY: + raise ValueError('forced f826 requires exact B5/N2944/D112/K512') + if key in EXACT_C829_KEYS: + return (_c829.launch_for_eval, dict(C829_SOURCE), 'guard_exact_two_k1024_keys_c829_v1', '(B, N, D, K) in {(4, 8192, 112, 1024), (4, 3712, 112, 1024)}') + source = source_for_shape(B=key[0], N=key[1], D=key[2], K=key[3]) + if key == EXACT_F826_KEY: + return (_retained.launch_for_eval_forced_f826, source, 'guard_exact_b5_n2944_d112_k512_f826_v1', 'B == 5 and N == 2944 and D == 112 and K == 512') + return (_retained.launch_for_eval, source, 'guard_assigned_d112_bucket_retained_4549_v1', 'assigned D112 bucket key excluding the c829 and f826 exact guards') + +def _source_value(source: dict[str, Any], suffix: str) -> Any: + for prefix in ('final_source', 'source', 'executable_body'): + value = source.get(''.join([format(prefix, ''), '_', format(suffix, '')])) + if value is not None: + return value + return None + +def _finish_outputs(inputs: dict[str, Any], child_outputs: dict[str, Any], source: dict[str, Any], guard_id: str, guard_condition: str) -> dict[str, Any]: + B, N, D, K = _shape(inputs) + key = (B, N, D, K) + caller_out = inputs['out'] + cluster_ids = child_outputs.get('cluster_ids', child_outputs.get('out')) + if cluster_ids is None: + raise ValueError(''.join([format(source['id'], ''), ' did not return cluster_ids or out'])) + caller_ptr = int(caller_out.data_ptr()) + returned_ptr = int(cluster_ids.data_ptr()) + if returned_ptr != caller_ptr: + raise ValueError(''.join([format(source['id'], ''), " replaced caller-owned inputs['out']"])) + child_trace = dict(child_outputs.get('route_trace') or {}) + specialized = key in EXACT_C829_KEYS or key == EXACT_F826_KEY + return {'cluster_ids': cluster_ids, 'selected_route': ROUTE_ID, 'route_trace': {'shape_key': _shape_key(key), 'selected_route': ROUTE_ID, 'selected_entrypoint': ''.join([format(__name__, ''), ':launch_for_eval']), 'selected_seed': _c829.SEED_ID if key in EXACT_C829_KEYS else child_trace.get('selected_seed'), 'portfolio_id': ROUTE_ID, 'portfolio_seed_id': SEED_ID, 'retained_portfolio': RETAINED_PORTFOLIO_SOURCE['id'], 'selected_source': source['id'], 'selected_source_task': source['source_task'], 'selected_source_revision': source.get('implementation_revision', source.get('implementation_commit', source.get('source_revision'))), 'selected_source_blob': _source_value(source, 'blob'), 'selected_source_sha256': _source_value(source, 'sha256'), 'selected_source_path': source['source_path'], 'selected_source_entrypoint': source['entrypoint'], 'child_selected_route': child_outputs.get('selected_route'), 'child_route_trace': child_trace, 'route_kind': 'specialized' if specialized else 'general', 'route_source': 'shape-specific-seed' if specialized else 'generated-variant', 'guard_id': guard_id, 'guard_condition': guard_condition, 'classification': 'seed-consumed' if specialized else 'route-ok', 'dispatch_key': 'B,N,D,K', 'dispatch_value': [B, N, D, K], 'caller_output_data_ptr': caller_ptr, 'returned_output_data_ptr': returned_ptr, 'caller_owned_output': True, 'output_dtype': 'int32', 'compute_kernel_count': 1, 'padding_count': 0, 'pack_count': 0, 'fallback_contract_regions': [], 'residual_contract_regions': [], 'global_workspace': False, 'materialized_padding': False, 'production_route': 'weave_only', 'launch_grid': child_trace.get('launch_grid'), 'producer_to_consumer': child_trace.get('producer_to_consumer'), 'primitive_family': child_trace.get('primitive_family', 'warp_mma_sync'), 'B': B, 'N': N, 'D': D, 'K': K}} + +def _launch(inputs: dict[str, Any], *, force_c829: bool=False, force_f826: bool=False) -> dict[str, Any]: + child, source, guard_id, guard_condition = _resolve_route(inputs, force_c829=force_c829, force_f826=force_f826) + return _finish_outputs(inputs, child(inputs), source, guard_id, guard_condition) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + """Launch exactly one frozen child through the caller-owned output ABI.""" + return _launch(inputs) + +def launch_for_eval_forced_c829(inputs: dict[str, Any]) -> dict[str, Any]: + return _launch(inputs, force_c829=True) + +def launch_for_eval_forced_f826(inputs: dict[str, Any]) -> dict[str, Any]: + return _launch(inputs, force_f826=True) + +@dataclass(frozen=True) +class PreparedPortfolioPlan: + inputs: dict[str, Any] + selected_route: str + selected_source: str + guard_id: str + direct_launcher: Callable[..., dict[str, Any]] + launch_count: int + device_index: int + stream: Any + stream_handle: int + caller_output_data_ptr: int + direct_launcher_token: str + + def launch(self, *, stream: Any=None, timeout_ms: float | None=None) -> dict[str, Any]: + import torch + with torch.cuda.device(self.device_index): + resolved_stream = self.stream if stream is None else stream + if int(resolved_stream.cuda_stream) != self.stream_handle: + raise RuntimeError('prepared a262 portfolio plan is stream-bound') + if int(self.inputs['out'].data_ptr()) != self.caller_output_data_ptr: + raise RuntimeError('prepared a262 caller output was replaced') + return self.direct_launcher(self.inputs, stream=None, timeout_ms=timeout_ms) + +def prepare_launch_plan(inputs: dict[str, Any], *, stream: Any=None, force_c829: bool=False, force_f826: bool=False) -> PreparedPortfolioPlan: + """Resolve and capture one stable direct child launch token.""" + import torch + from .._dispatch_runtime import detect_gpu_arch + from .._dispatch_runtime import capture_kernel_launches + device_index = inputs['x'].device.index + if device_index is None: + device_index = torch.cuda.current_device() + device_index = int(device_index) + with torch.cuda.device(device_index): + resolved_stream = torch.cuda.current_stream(device_index) if stream is None else stream + stream_device = getattr(resolved_stream, 'device', None) + stream_device_index = getattr(stream_device, 'index', stream_device) + if stream_device_index is not None and int(stream_device_index) != device_index: + raise ValueError('prepared a262 stream device does not match inputs') + if detect_gpu_arch() not in {'sm_100a', 'sm_103a'}: + raise ValueError('a262 prepared routes require sm_100a or sm_103a') + child, source, guard_id, guard_condition = _resolve_route(inputs, force_c829=force_c829, force_f826=force_f826) + with _PREPARE_LOCK: + with capture_kernel_launches(stream=resolved_stream) as captured: + child_outputs = child(inputs) + prepared_result = _finish_outputs(inputs, child_outputs, source, guard_id, guard_condition) + direct_launcher = captured.bind(prepared_result) + return PreparedPortfolioPlan(inputs=inputs, selected_route=ROUTE_ID, selected_source=source['id'], guard_id=guard_id, direct_launcher=direct_launcher, launch_count=direct_launcher.launch_count, device_index=device_index, stream=resolved_stream, stream_handle=int(resolved_stream.cuda_stream), caller_output_data_ptr=int(inputs['out'].data_ptr()), direct_launcher_token=''.join([format(type(direct_launcher).__module__, ''), '.', format(type(direct_launcher).__qualname__, ''), '@', format(id(direct_launcher), ''.join(['x']))])) + +def launch_prepared(plan: PreparedPortfolioPlan, *, stream: Any=None, timeout_ms: float | None=None) -> dict[str, Any]: + if not isinstance(plan, PreparedPortfolioPlan): + raise TypeError('plan must be returned by prepare_launch_plan') + return plan.launch(stream=stream, timeout_ms=timeout_ms) diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_f826_exact_portfolio_generalize_auto_tuning_flash_kmeans_assign_8d82_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_f826_exact_portfolio_generalize_auto_tuning_flash_kmeans_assign_8d82_v1.py new file mode 100644 index 00000000..e02905a6 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_f826_exact_portfolio_generalize_auto_tuning_flash_kmeans_assign_8d82_v1.py @@ -0,0 +1,134 @@ +"""Namespaced D112 portfolio with one exact f826 guard (minimum: sm_100a). + +The guard ``B=5, N=2944, D=112, K=512`` selects the validated f826 +two-owner warp-MMA seed. K=1024 remains on the frozen c72b overlay and every +other admitted D112 row remains on frozen 4549. This module is additive: it +does not alter the production dispatcher or any child kernel schedule. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from dataclasses import dataclass +from threading import RLock +from typing import Any, Callable +from . import flash_kmeans_assign_d112_4549_c72b_k1024_overlay_weave_evolve_flash_kmeans_assign_3ae6_v1 as _frozen +from . import flash_kmeans_assign_d112_k512_two_owner_peer_only_mma_weave_evolve_flash_kmeans_assign_f826_v1 as _f826 +FEAT_D = 112 +EXACT_F826_KEY = (5, 2944, 112, 512) +SUPPORTED_K = _frozen.SUPPORTED_K +ROUTE_ID = 'd112_f826_exact_portfolio_generalize_auto_tuning_flash_kmeans_assign_8d82_v1' +SEED_ID = 'd112-overlay-plus-f826-exact-k512-consumption-v1' +TARGET_SHAPE = 'physical_D112_padded_14' +EXACT_SHAPE_KEY = 'adjacent_68cf_d112_tail_b5_n2944_k512_d112' +EXACT_GUARD_ID = 'guard_exact_b5_n2944_d112_k512_f826_v1' +EXACT_GUARD_CONDITION = 'B == 5 and N == 2944 and D == 112 and K == 512' +F826_SOURCE = {'id': 'f826_exact_k512', 'source_task': 'weave-evolve-flash-kmeans-assign-f826', 'implementation_commit': '124ef2ad24c3ac86c063e69da0cb0c5d7fca0174', 'implementation_source_blob': '3a1dee841855c15481a740f28ba07142e4fdbb29', 'implementation_source_sha256': '99c44085f258a699b7e3372115cc82aa969643addae0797bf47c0758cad8cf0f', 'measurement_source_blob': 'a3dcf64c5014d223127bf0a39b13c485cf140c48', 'measurement_source_sha256': 'fd3e97727367d3acbbd3b994a86cb25f563ce3f32544a724dbd09f05a70a88ce', 'final_source_blob': 'f30b546e70bfbc80a94547458d857ead0de5e320', 'final_source_sha256': '47c2d22bbb00b3061d268fae9c405e62257bde8ba43ba33fc1124d456572cd63', 'source_path': 'loom/examples/weave/flash_kmeans_assign_d112_k512_two_owner_peer_only_mma_weave_evolve_flash_kmeans_assign_f826_v1.py', 'entrypoint': ''.join([format(_f826.__name__, ''), ':launch_for_eval'])} +_PREPARE_LOCK = RLock() + +def _shape(inputs: dict[str, Any]) -> tuple[int, int, int, int]: + return tuple((int(inputs[key]) for key in ('B', 'N', 'D', 'K'))) + +def _is_exact_f826(*, B: int, N: int, D: int, K: int) -> bool: + return (B, N, D, K) == EXACT_F826_KEY + +def supports_shape(*, D: int, N: int, K: int, B: int | None=None) -> bool: + del B + return _frozen.supports_shape(D=D, N=N, K=K) + +def source_for_shape(*, B: int, N: int, D: int, K: int) -> dict[str, Any]: + if _is_exact_f826(B=B, N=N, D=D, K=K): + return dict(F826_SOURCE) + return dict(_frozen.source_for_k(K)) + +def _resolve_route(inputs: dict[str, Any], *, force_f826: bool=False) -> tuple[Callable[[dict[str, Any]], dict[str, Any]], dict[str, Any], str, str]: + B, N, D, K = _shape(inputs) + if not supports_shape(B=B, D=D, N=N, K=K): + raise ValueError('8d82 D112 portfolio requires D=112, N%64=0, and K in (256, 512, 768, 1024, 4096, 8192)') + exact = _is_exact_f826(B=B, N=N, D=D, K=K) + if force_f826 and (not exact): + raise ValueError('forced f826 route requires exact B5/N2944/D112/K512') + if exact: + return (_f826.launch_for_eval, dict(F826_SOURCE), EXACT_GUARD_ID, EXACT_GUARD_CONDITION) + source = dict(_frozen.source_for_k(K)) + return (_frozen.launch_for_eval, source, 'guard_frozen_3ae6_k_ownership', 'K == 1024 selects c72b; other admitted K values select 4549') + +def _finish_outputs(inputs: dict[str, Any], child_outputs: dict[str, Any], source: dict[str, Any], guard_id: str, guard_condition: str) -> dict[str, Any]: + B, N, D, K = _shape(inputs) + caller_out = inputs['out'] + cluster_ids = child_outputs.get('cluster_ids', child_outputs.get('out')) + if cluster_ids is None: + raise ValueError(''.join([format(source['id'], ''), ' did not return cluster_ids or out'])) + caller_ptr = int(caller_out.data_ptr()) + returned_ptr = int(cluster_ids.data_ptr()) + if returned_ptr != caller_ptr: + raise ValueError(''.join([format(source['id'], ''), " did not retain caller-owned inputs['out']"])) + child_trace = dict(child_outputs.get('route_trace') or {}) + exact = source['id'] == F826_SOURCE['id'] + return {'cluster_ids': cluster_ids, 'selected_route': ROUTE_ID, 'route_trace': {'shape_key': EXACT_SHAPE_KEY if exact else TARGET_SHAPE, 'selected_route': ROUTE_ID, 'selected_entrypoint': ''.join([format(__name__, ''), ':launch_for_eval']), 'selected_seed': _f826.SEED_ID if exact else child_trace.get('selected_seed'), 'expected_seed': _f826.SEED_ID if exact else child_trace.get('selected_seed'), 'portfolio_id': ROUTE_ID, 'portfolio_seed_id': SEED_ID, 'selected_source': source['id'], 'selected_source_task': source['source_task'], 'selected_source_revision': source.get('implementation_commit', source.get('source_revision')), 'selected_source_blob': source.get('final_source_blob', source.get('source_blob')), 'selected_source_sha256': source.get('final_source_sha256', source.get('source_sha256')), 'selected_source_path': source['source_path'], 'selected_source_entrypoint': source['entrypoint'], 'child_selected_route': child_outputs.get('selected_route'), 'child_route_trace': child_trace, 'route_kind': 'specialized' if exact else 'general', 'route_source': 'shape-specific-seed' if exact else 'generated-variant', 'guard_id': guard_id, 'guard_condition': guard_condition, 'classification': 'seed-consumed' if exact else 'route-ok', 'dispatch_key': 'B,N,D,K' if exact else 'K', 'dispatch_value': [B, N, D, K] if exact else K, 'caller_output_data_ptr': caller_ptr, 'returned_output_data_ptr': returned_ptr, 'caller_owned_output': True, 'output_dtype': 'int32', 'compute_kernel_count': 1, 'padding_count': 0, 'pack_count': 0, 'fallback_contract_regions': [], 'residual_contract_regions': [], 'global_workspace': False, 'materialized_padding': False, 'production_route': 'weave_only', 'launch_grid': child_trace.get('launch_grid'), 'producer_to_consumer': child_trace.get('producer_to_consumer'), 'B': B, 'N': N, 'D': D, 'K': K}} + +def _launch(inputs: dict[str, Any], *, force_f826: bool=False) -> dict[str, Any]: + child, source, guard_id, guard_condition = _resolve_route(inputs, force_f826=force_f826) + return _finish_outputs(inputs, child(inputs), source, guard_id, guard_condition) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + """Launch the guarded namespaced portfolio through its public ABI.""" + return _launch(inputs) + +def launch_for_eval_forced_f826(inputs: dict[str, Any]) -> dict[str, Any]: + """Force the exact f826 route for exported-ABI correctness evidence.""" + return _launch(inputs, force_f826=True) + +@dataclass(frozen=True) +class PreparedPortfolioPlan: + inputs: dict[str, Any] + selected_route: str + selected_source: str + guard_id: str + direct_launcher: Callable[..., dict[str, Any]] + launch_count: int + device_index: int + stream: Any + stream_handle: int + caller_output_data_ptr: int + direct_launcher_token: str + + def launch(self, *, stream: Any=None, timeout_ms: float | None=None) -> dict[str, Any]: + """Submit the captured child launch without re-entering either dispatcher.""" + import torch + with torch.cuda.device(self.device_index): + resolved_stream = self.stream if stream is None else stream + if int(resolved_stream.cuda_stream) != self.stream_handle: + raise RuntimeError('prepared 8d82 portfolio plan is stream-bound') + if int(self.inputs['out'].data_ptr()) != self.caller_output_data_ptr: + raise RuntimeError('prepared 8d82 portfolio caller output was replaced') + return self.direct_launcher(self.inputs, stream=None, timeout_ms=timeout_ms) + +def prepare_launch_plan(inputs: dict[str, Any], *, stream: Any=None, force_f826: bool=False) -> PreparedPortfolioPlan: + """Resolve the exact route once and capture one stable direct launch token.""" + import torch + from .._dispatch_runtime import detect_gpu_arch + from .._dispatch_runtime import capture_kernel_launches + device_index = inputs['x'].device.index + if device_index is None: + device_index = torch.cuda.current_device() + device_index = int(device_index) + with torch.cuda.device(device_index): + resolved_stream = torch.cuda.current_stream(device_index) if stream is None else stream + stream_device = getattr(resolved_stream, 'device', None) + stream_device_index = getattr(stream_device, 'index', stream_device) + if stream_device_index is not None and int(stream_device_index) != device_index: + raise ValueError('prepared 8d82 stream device does not match input device') + if detect_gpu_arch() not in {'sm_100a', 'sm_103a'}: + raise ValueError('8d82 prepared routes require sm_100a or sm_103a') + child, source, guard_id, guard_condition = _resolve_route(inputs, force_f826=force_f826) + with _PREPARE_LOCK: + with capture_kernel_launches(stream=resolved_stream) as captured: + child_outputs = child(inputs) + prepared_result = _finish_outputs(inputs, child_outputs, source, guard_id, guard_condition) + direct_launcher = captured.bind(prepared_result) + return PreparedPortfolioPlan(inputs=inputs, selected_route=ROUTE_ID, selected_source=source['id'], guard_id=guard_id, direct_launcher=direct_launcher, launch_count=direct_launcher.launch_count, device_index=device_index, stream=resolved_stream, stream_handle=int(resolved_stream.cuda_stream), caller_output_data_ptr=int(inputs['out'].data_ptr()), direct_launcher_token=''.join([format(type(direct_launcher).__module__, ''), '.', format(type(direct_launcher).__qualname__, ''), '@', format(id(direct_launcher), ''.join(['x']))])) + +def launch_prepared(plan: PreparedPortfolioPlan, *, stream: Any=None, timeout_ms: float | None=None) -> dict[str, Any]: + if not isinstance(plan, PreparedPortfolioPlan): + raise TypeError('plan must be returned by prepare_launch_plan') + return plan.launch(stream=stream, timeout_ms=timeout_ms) diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_k1024_owner_local_warp_mma_distinct_workfeed_weave_evolve_flash_kmeans_assign_c829_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_k1024_owner_local_warp_mma_distinct_workfeed_weave_evolve_flash_kmeans_assign_c829_v1.py new file mode 100644 index 00000000..6c47e4ad --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_k1024_owner_local_warp_mma_distinct_workfeed_weave_evolve_flash_kmeans_assign_c829_v1.py @@ -0,0 +1,68 @@ +"""D112 K1024 owner-local streamed warp-MMA path (minimum arch: sm_90a). + +One ordinary CTA owns one N64 point tile. Warp 0 loads the point tile once +and streams alternating K32 centroid slabs. Four row-owner warps retain the +same 16 point rows for the full K1024 scan, consume each slab with +``ldmatrix`` + ``mma.sync``, and write caller-owned cluster IDs directly. + +This physical DAG has no K-sharded CTA owners, cluster/DSM handoff, TMEM, +global workspace, padding, pack kernel, or post-kernel reduction. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +BLOCK_N = 64 +BLOCK_K = 32 +FEAT_D = 112 +FEATURE_TILES = FEAT_D // 16 +LOAD_WARPS = 1 +COMPUTE_WARPS = 4 +THREADS = (LOAD_WARPS + COMPUTE_WARPS) * 32 +K_TILES = 1024 // BLOCK_K +K_TILE_PAIRS = K_TILES // 2 +X_ELEMS = BLOCK_N * FEAT_D +C_ELEMS = BLOCK_K * FEAT_D +X_BYTES = X_ELEMS * 2 +C_BYTES = C_ELEMS * 2 +X_OFFSET = 0 +C0_OFFSET = X_OFFSET + X_BYTES +C1_OFFSET = C0_OFFSET + C_BYTES +SMEM_BYTES = C1_OFFSET + C_BYTES +ROUTE_ID = 'd112_k1024_owner_local_warp_mma_distinct_workfeed_weave_evolve_flash_kmeans_assign_c829_v1' +SEED_ID = 'd112-k1024-owner-local-warp-mma-distinct-workfeed-weave-evolve-flash-kmeans-assign-c829-v1' +TARGET_SHAPE = 'post_d895_d112_b4_n8192_k1024_d112' +CAPABILITY_SHAPE_KEYS = (TARGET_SHAPE, 'adjacent_ef00_d112_odd_tail_b4_n3712_k1024_d112') +flash_kmeans_assign_d112_k1024_owner_local_warp_mma_distinct_workfeed_c829_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d112_k1024_owner_local_warp_mma_distinct_workfeed_c829_v1", "arg_keys": ["x", "centroids", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 29696, "constants": [], "cta_group": 1, "threads": 160}')) +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d112_k1024_owner_local_warp_mma_distinct_workfeed_c829_v1", "arg_keys": ["x", "centroids", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 29696, "constants": [], "cta_group": 1, "threads": 160}')) + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as include_dirs + return include_dirs() + +def _compiled() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0353"}, "kernel_flash_kmeans_assign_d112_k1024_owner_local_warp_mma_distinct_workfeed_c829_v1", 29696, 160]}')) + +def supports_shape(*, B: int, D: int, N: int, K: int) -> bool: + return B == 4 and D == FEAT_D and (N in (8192, 3712)) and (K == 1024) + +def _shape_key(*, B: int, N: int, D: int, K: int) -> str: + if (B, N, D, K) == (4, 8192, 112, 1024): + return CAPABILITY_SHAPE_KEYS[0] + if (B, N, D, K) == (4, 3712, 112, 1024): + return CAPABILITY_SHAPE_KEYS[1] + raise ValueError(''.join(['unassigned D112 K1024 capability shape ', format(repr((B, N, D, K)), '')])) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + import torch + from .._dispatch_runtime import CUDAKernel + B, N, D, K = (int(inputs[key]) for key in ('B', 'N', 'D', 'K')) + if not supports_shape(B=B, D=D, N=N, K=K): + raise ValueError('c829 owner-local warp-MMA seed admits only B4/D112/K1024 with N in (8192, 3712)') + cubin, name, smem, threads = _compiled() + grid = (B * (N // BLOCK_N), 1, 1) + with CUDAKernel(cubin, name) as kernel: + kernel.launch(grid=grid, block=(threads, 1, 1), args=[inputs['x'], inputs['centroids'], inputs['c_sq'], inputs['out'], B, N, D, K, N // BLOCK_N], shared_mem=smem, stream=torch.cuda.current_stream(), timeout_ms=120000) + shape_key = _shape_key(B=B, N=N, D=D, K=K) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': {'shape_key': shape_key, 'selected_route': ROUTE_ID, 'selected_entrypoint': ''.join([format(__name__, ''), ':launch_for_eval']), 'selected_seed': SEED_ID, 'route_kind': 'single_cta_owner_local_k32_streamed_warp_mma', 'primitive_family': 'warp_mma_sync', 'compute_kernel_count': 1, 'cluster_ctas': 1, 'cta_group': 1, 'point_tile': BLOCK_N, 'centroid_tile': BLOCK_K, 'loader_warps': LOAD_WARPS, 'row_owner_warps': COMPUTE_WARPS, 'producer_ownership': 'warp 0 stages one shared point tile and alternating full-K centroid slabs', 'consumer_ownership': 'four row-owner warps retain disjoint 16-row spans for all K1024', 'producer_to_consumer': 'bulk_gmem_to_double_buffered_smem_to_ldmatrix_to_mma_sync_to_owner_local_argmax_to_caller_cluster_ids', 'split_k_reduction': 'none', 'global_workspace': False, 'padding_count': 0, 'pack_count': 0, 'fallback_contract_regions': [], 'residual_contract_regions': [], 'launch_grid': grid, 'num_n_tiles': N // BLOCK_N, 'shared_pool_bytes': SMEM_BYTES, 'shared_storage_bytes': int(ir.computed_smem_bytes)}} diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_k512_two_owner_peer_only_mma_weave_evolve_flash_kmeans_assign_f826_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_k512_two_owner_peer_only_mma_weave_evolve_flash_kmeans_assign_f826_v1.py new file mode 100644 index 00000000..d455e7c2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_k512_two_owner_peer_only_mma_weave_evolve_flash_kmeans_assign_f826_v1.py @@ -0,0 +1,66 @@ +"""D112 K512 two-owner peer-only warp-MMA path (minimum architecture: sm_90a). + +Two independent ``cta_group=1`` owners split K512 into disjoint K256 spans. +Each owner uses two four-warp consumer groups with direct TMA centroid stages, +``ldmatrix.x2``, and ``mma.sync`` before locally merging packed argmax keys. +Rank 1 sends one 512-byte key slab to rank 0; rank 0 retains its local slab and +writes caller-owned ``int32`` cluster IDs without a self-copy or workspace. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +BLOCK_N, BLOCK_K, FEAT_D, THREADS = (64, 8, 112, 256) +CLUSTER_CTAS, CONSUMER_GROUPS, PIPELINE_STAGES = (2, 2, 2) +SUPPORTED_K = (512,) +X_ELEMS, C_ELEMS = (BLOCK_N * FEAT_D, BLOCK_K * FEAT_D) +X_BYTES, C_BYTES = (X_ELEMS * 2, C_ELEMS * 2) +CENTROID_STAGE_COUNT = CONSUMER_GROUPS * PIPELINE_STAGES +SCORE_BYTES = CONSUMER_GROUPS * 4 * 16 * BLOCK_K * 4 +LOCAL_KEY_BYTES = BLOCK_N * 8 +GROUP_KEY_BYTES = CONSUMER_GROUPS * LOCAL_KEY_BYTES +PEER_KEY_BYTES = LOCAL_KEY_BYTES +U32_MASK = 4294967295 +X_RAW_OFFSET = 0 +C_DIRECT_OFFSET = X_RAW_OFFSET + X_BYTES +X_MMA_OFFSET = C_DIRECT_OFFSET + CENTROID_STAGE_COUNT * C_BYTES +SCORE_OFFSET = X_MMA_OFFSET + X_BYTES +GROUP_KEY_OFFSET = SCORE_OFFSET + SCORE_BYTES +LOCAL_KEY_OFFSET = GROUP_KEY_OFFSET + GROUP_KEY_BYTES +PEER_KEY_OFFSET = LOCAL_KEY_OFFSET + LOCAL_KEY_BYTES +SMEM_BYTES = PEER_KEY_OFFSET + PEER_KEY_BYTES +ROUTE_ID = 'd112_k512_two_owner_peer_only_mma_weave_evolve_flash_kmeans_assign_f826_v1' +SEED_ID = 'd112-k512-two-owner-peer-only-mma-weave-evolve-flash-kmeans-assign-f826-v1' +TARGET_SHAPE = 'adjacent_68cf_d112_tail_b5_n2944_k512_d112' +flash_kmeans_assign_d112_k512_two_owner_peer_only_mma_f826_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d112_k512_two_owner_peer_only_mma_f826_v1", "arg_keys": ["x", "centroids", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles"], "cluster_dims": [2, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 256}')) +cluster_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d112_k512_two_owner_peer_only_mma_f826_v1", "arg_keys": ["x", "centroids", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles"], "cluster_dims": [2, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 256}')) +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d112_k512_two_owner_peer_only_mma_f826_v1", "arg_keys": ["x", "centroids", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles"], "cluster_dims": [2, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 256}')) + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as dirs + return dirs() + +def _compiled_cluster() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0354"}, "kernel_flash_kmeans_assign_d112_k512_two_owner_peer_only_mma_f826_v1", 43008, 256]}')) + +def supports_shape(*, D: int, N: int, K: int) -> bool: + return D == FEAT_D and N % BLOCK_N == 0 and (K in SUPPORTED_K) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + import torch + from .._dispatch_runtime import CUDAKernel + B, N, D, K = (int(inputs[key]) for key in ('B', 'N', 'D', 'K')) + if not supports_shape(D=D, N=N, K=K): + raise ValueError('two-owner peer-only MMA f826 requires D=112, N%64=0, and K=512') + caller_out = inputs['out'] + caller_ptr = int(caller_out.data_ptr()) + cubin, name, smem, threads = _compiled_cluster() + num_n_tiles = N // BLOCK_N + grid = (B * num_n_tiles * CLUSTER_CTAS, 1, 1) + with CUDAKernel(cubin, name) as kernel: + kernel.launch_cluster(grid=grid, block=(threads, 1, 1), args=[inputs['x'], inputs['centroids'], inputs['c_sq'], caller_out, B, N, D, K, num_n_tiles], cluster_dims=(CLUSTER_CTAS, 1, 1), shared_mem=smem, stream=torch.cuda.current_stream(), timeout_ms=120000) + returned_ptr = int(caller_out.data_ptr()) + if returned_ptr != caller_ptr: + raise ValueError('two-owner peer-only MMA did not retain caller-owned out') + return {'cluster_ids': caller_out, 'selected_route': ROUTE_ID, 'route_trace': {'shape_key': TARGET_SHAPE, 'selected_route': ROUTE_ID, 'selected_entrypoint': ''.join([format(__name__, ''), ':launch_for_eval']), 'selected_seed': SEED_ID, 'route_kind': 'd112_k512_two_owner_peer_only_warp_mma', 'cluster_ctas': CLUSTER_CTAS, 'cta_group': 1, 'consumer_groups_per_owner': CONSUMER_GROUPS, 'centroid_pipeline_stages': PIPELINE_STAGES, 'centroids_per_owner': K // CLUSTER_CTAS, 'centroids_per_consumer_group': K // (CLUSTER_CTAS * CONSUMER_GROUPS), 'producer_ownership': 'ranks 0..1 independently own K256; each four-warp group owns two direct canonical centroid stages', 'reuse_mechanism': 'direct TMA canonical stages consumed by ldmatrix.x2 with independent 128-thread group barriers', 'producer_to_consumer': 'D112_TMA_to_ldmatrix_x2_to_mma_sync_to_local_two_group_argmax_to_one_peer_key_handoff_to_caller_cluster_ids', 'split_k_reduction': 'per-owner shared two-group key merge; one rank1-to-rank0 peer slab; rank0 retains local slab', 'remote_producer_count': 1, 'peer_key_transaction_bytes': PEER_KEY_BYTES, 'rank0_self_copy': False, 'caller_output_data_ptr': caller_ptr, 'returned_output_data_ptr': returned_ptr, 'caller_owned_output': True, 'output_dtype': 'int32', 'compute_kernel_count': 1, 'padding_count': 0, 'pack_count': 0, 'global_workspace': False, 'fallback_contract_regions': [], 'residual_contract_regions': [], 'launch_grid': grid, 'B': B, 'N': N, 'D': D, 'K': K, 'num_n_tiles': num_n_tiles}} diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_shared_point_dual_issuer_tcgen05_weave_evolve_flash_kmeans_assign_6e1e_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_shared_point_dual_issuer_tcgen05_weave_evolve_flash_kmeans_assign_6e1e_v1.py new file mode 100644 index 00000000..4100bae3 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d112_shared_point_dual_issuer_tcgen05_weave_evolve_flash_kmeans_assign_6e1e_v1.py @@ -0,0 +1,87 @@ +"""D112 shared-point dual-independent-issuer tcgen05/TMEM assignment. + +Minimum architecture: sm_100a. A dedicated point producer stages the seven +native D16 slices of each N128 tile exactly once. Two independently +synchronized centroid-owner pipelines consume that shared point tile through +distinct tcgen05 issuer warps and disjoint K64 TMEM score slots, merge ordered +argmax keys on chip, and write the caller-owned ``cluster_ids`` tensor. Every +issuer owns its operand-ready, score-complete, readback-complete, and recycle +phase state. The production launch is Weave-only: there is no D128 adapter, +DSM point fanout, global workspace, or external fallback. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +BLOCK_N, BLOCK_K, FEAT_D, FEAT_CHUNK = (128, 64, 112, 16) +FEATURE_TILES = FEAT_D // FEAT_CHUNK +OWNER_SLOTS, COMPUTE_WARPS = (2, 4) +THREADS = 13 * 32 +X_SLICE_BYTES = BLOCK_N * FEAT_CHUNK * 2 +C_SLICE_BYTES = BLOCK_K * FEAT_CHUNK * 2 +LOCAL_KEY_BYTES = BLOCK_N * 8 +X_BYTES = FEATURE_TILES * X_SLICE_BYTES +C0_OFFSET = X_BYTES +C1_OFFSET = C0_OFFSET + C_SLICE_BYTES +KEY0_OFFSET = C1_OFFSET + C_SLICE_BYTES +KEY1_OFFSET = KEY0_OFFSET + LOCAL_KEY_BYTES +SMEM_BYTES = KEY1_OFFSET + LOCAL_KEY_BYTES +U32_MASK = 4294967295 +TCGEN05_ARCHES = frozenset({'sm_100a', 'sm_103a'}) +ROUTE_ID = 'd112_shared_point_dual_issuer_tcgen05_weave_evolve_flash_kmeans_assign_6e1e_v1' +SEED_ID = 'd112-shared-point-dual-issuer-tcgen05-weave-evolve-flash-kmeans-assign-6e1e-v1' +TARGET_SHAPE = 'physical_D112_padded_14' +CAPABILITY_SHAPE = 'post_d895_d112_b4_n8192_k1024_d112' +PHYSICAL_DAG = 'one_N128_D112_point_producer->seven_persistent_D16_slices->two_independent_K64_tcgen05_issuers->two_disjoint_TMEM_argmax_consumers->local_packed_key_merge->caller_cluster_ids' +flash_kmeans_assign_d112_shared_point_dual_issuer_tcgen05_6e1e_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d112_shared_point_dual_issuer_tcgen05_6e1e_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 35840, "constants": [], "cta_group": 1, "threads": 416}')) +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d112_shared_point_dual_issuer_tcgen05_6e1e_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 35840, "constants": [], "cta_group": 1, "threads": 416}')) + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as include_dirs + return include_dirs() + +def _require_tcgen05_arch() -> str: + from .._dispatch_runtime import detect_gpu_arch + arch = detect_gpu_arch() + if arch not in TCGEN05_ARCHES: + supported = ', '.join(sorted(TCGEN05_ARCHES)) + raise RuntimeError(''.join([format(ROUTE_ID, ''), ' requires tcgen05/TMEM (', format(supported, ''), '); detected ', format(arch, '')])) + return arch + +def _compiled() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0355"}, "kernel_flash_kmeans_assign_d112_shared_point_dual_issuer_tcgen05_6e1e_v1", 35840, 416]}')) + +def supports_shape(*, D: int, N: int, K: int) -> bool: + return D == FEAT_D and N % BLOCK_N == 0 and (K % (BLOCK_K * OWNER_SLOTS) == 0) + +def _make_tmaps(inputs: dict[str, Any]) -> tuple[Any, Any]: + from .._dispatch_runtime import create_tensor_map_3d_32b + bsz, n_points, n_clusters = (int(inputs[key]) for key in ('B', 'N', 'K')) + return (create_tensor_map_3d_32b(inputs['x'].data_ptr(), bsz * n_points, BLOCK_N, FEAT_D, FEAT_CHUNK), create_tensor_map_3d_32b(inputs['centroids'].data_ptr(), bsz * n_clusters, BLOCK_K, FEAT_D, FEAT_CHUNK)) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + import torch + from .._dispatch_runtime import CUDAKernel + _require_tcgen05_arch() + B, N, D, K = (int(inputs[key]) for key in ('B', 'N', 'D', 'K')) + if not supports_shape(D=D, N=N, K=K): + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires D=112, N%128=0, and K%128=0'])) + x_tmap, c_tmap = _make_tmaps(inputs) + cubin, name, smem, threads = _compiled() + grid = (B * (N // BLOCK_N), 1, 1) + with CUDAKernel(cubin, name) as kernel: + kernel.launch(grid=grid, block=(threads, 1, 1), args=[x_tmap, c_tmap, inputs['c_sq'], inputs['out'], B, N, D, K, N // BLOCK_N, K // BLOCK_K], shared_mem=smem, stream=torch.cuda.current_stream()) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': {'shape_key': TARGET_SHAPE, 'capability_shape': CAPABILITY_SHAPE, 'selected_route': ROUTE_ID, 'selected_entrypoint': ''.join([format(__name__, ''), ':launch_for_eval']), 'selected_seed': SEED_ID, 'primitive': 'native_D112_TMA+tcgen05_mma+TMEM', 'physical_dag': PHYSICAL_DAG, 'point_producer': 'warp0 loads seven persistent D16 slices once per N128 tile', 'centroid_ownership': 'slot0 owns even K64 tiles; slot1 owns odd K64 tiles', 'issuer_policy': 'two independent warp issuers with no shared completion barrier', 'issuer_ownership': [{'slot': 0, 'issuer_warp': 3, 'operand_ready': ['x_ready0', 'c_full0'], 'score_complete': 'score_full0', 'readback_complete': 'score_empty0', 'recycle': ['c_empty0', 'x_empty0']}, {'slot': 1, 'issuer_warp': 12, 'operand_ready': ['x_ready1', 'c_full1'], 'score_complete': 'score_full1', 'readback_complete': 'score_empty1', 'recycle': ['c_empty1', 'x_empty1']}], 'score_slots': [{'tmem_cols': [0, 64], 'issuer_warp': 3, 'compute_warps': [4, 5, 6, 7]}, {'tmem_cols': [64, 128], 'issuer_warp': 12, 'compute_warps': [8, 9, 10, 11]}], 'producer_to_consumer': 'shared_seven_D16_point_slices_to_two_independent_K64_tcgen05_issuers_to_disjoint_TMEM_argmax_consumers_to_local_ordered_key_merge_to_cluster_ids', 'output_ownership': 'caller_owned_inputs[out]', 'production_launches': [''.join(['kernel_', format(ir.symbol, '')])], 'fallback_contract_regions': [], 'residual_contract_regions': [], 'global_workspace': False, 'materialized_d128_adapter': False, 'launch_grid': grid}} + +def benchmark_shared_point_dual_issuer_tcgen05_6e1e(*, shape_label: str=CAPABILITY_SHAPE, benchmark: bool=True) -> dict[str, Any]: + """Registered one-row capability launcher for regression and sanitizer tools.""" + from .. import _dispatch_runtime as padded + shape = next((row for row in padded.padded_bucket_shapes('d112_gap_pad_14') if row['label'] == shape_label)) + report = padded.evaluate(launch_for_eval, shapes=[shape], correctness=True, benchmark=benchmark, flashlib_baseline=False, benchmark_warmup_ms=20.0, benchmark_ms=100.0) + row = report['per_shape'][shape_label] + exact = bool(report['correctness']['all_correct']) and row.get('cluster_id_exact_match') is True + result: dict[str, Any] = {'passed': exact, 'all_pass': exact, 'shape_label': shape_label, 'selected_route': row.get('selected_route'), 'cluster_id_exact_match': row.get('cluster_id_exact_match'), 'measurement_comparable': bool(row.get('measurement_comparable')) if benchmark else False} + if benchmark: + result.update({'tflops': float(row['tflops']), 'kernel_ms': float(row['kernel_ms']), 'timing_backend': row.get('timing_backend'), 'timing_backend_requested': row.get('timing_backend_requested'), 'timing_backend_fallback_reason': row.get('timing_backend_fallback_reason'), 'performance_comparable': bool(row.get('measurement_comparable')), 'measured_entrypoint': ''.join([format(__name__, ''), ':launch_for_eval'])}) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d128_splitk_priority_575c_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d128_splitk_priority_575c_v1.py new file mode 100644 index 00000000..bae39528 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d128_splitk_priority_575c_v1.py @@ -0,0 +1,37 @@ +"""Exact-D128 no-padding Split-K Flash-KMeans priority seed. + +Minimum architecture: sm_100a. This D=128-only seed reuses the validated +BLOCK_N=64 G1/R4 tcgen05/TMEM producer and four-lane reducer, but exposes every +256-centroid tile as a Split-K work item for the five priority K=4096/8192 +rows. It must not be used on sm_120a/sm_121a, where ptxas rejects tcgen05. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from typing import Any +from . import flash_kmeans_assign_highd_splitk_blockn64_g1r4_streamdep_r63_v1 as _g1r4 +ROUTE_ID = 'd128_splitk_priority_575c_v1' +SEED_ID = 'd128-no-padding-splitk-priority-575c-v1' +TARGET_D = 128 +MIN_K_TILES = 16 +_g1r4.SUPPORTED_DIMS.add(TARGET_D) + +def _use_d128_priority_splitk(*, dim: int, num_n_tiles: int, k_tiles: int) -> bool: + return dim == TARGET_D and num_n_tiles <= _g1r4.MAX_POINT_TILES and (k_tiles >= MIN_K_TILES) +_g1r4._use_blockn64_splitk = _use_d128_priority_splitk +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_partial_blockn64_g1r4_streamdep_r63_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 192}')) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + """Launch only the exact D128, no-padding priority Split-K schedule.""" + if int(inputs['B']) != 1: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires B=1'])) + if int(inputs['D']) != TARGET_D: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires D=', format(TARGET_D, ''), ', got ', format(inputs['D'], '')])) + if int(inputs['N']) not in (512, 1024, 2048): + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires N in {512, 1024, 2048}'])) + if int(inputs['K']) not in (4096, 8192): + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires K in {4096, 8192}'])) + outputs = _g1r4.launch_for_eval(inputs) + trace = dict(outputs.get('route_trace', {})) + trace.update({'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_d128_splitk_priority_575c_v1:launch_for_eval', 'selected_seed': SEED_ID, 'guard_id': 'guard_d128_no_padding_splitk_priority_575c_v1', 'guard_condition': 'B == 1 and D == 128 and N in [512,1024,2048] and K in [4096,8192]', 'reason': 'exact-D128 priority Split-K: one 256-K tile per G1 producer work item plus four-lane reducer'}) + return {'cluster_ids': outputs['cluster_ids'], 'selected_route': ROUTE_ID, 'route_trace': trace} diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1.py new file mode 100644 index 00000000..00524c33 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1.py @@ -0,0 +1,141 @@ +"""Flash-KMeans Euclidean assignment D144/D160/D176 pad-to-D192 seed. + +Minimum architecture: sm_100a. This candidate zero-pads non-128D BF16 inputs +with a Weave pack kernel, then routes assignment to the existing D192 +tcgen05/TMEM split-D seed. It is not intended for sm_120a/sm_121a where ptxas +rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +from . import flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1 as _d192 +from .flash_kmeans_assign_stream import stream_cache_key +BLOCK_N = _d192.BLOCK_N +BLOCK_K = _d192.BLOCK_K +FEAT_D_PAD = 192 +FEAT_D_PAD_VECS = FEAT_D_PAD // 8 +SUPPORTED_D = {144, 160, 176} +PACK_THREADS = 256 +PACK_GRID_CAP = 4096 +_SCRATCH_CACHE: dict[tuple[int, ...], tuple[Any, Any]] = {} +_ROUTE_INPUT_CACHE: dict[tuple[int, ...], dict[str, Any]] = {} +flash_kmeans_assign_d160_pad192_pack_f9b2_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d160_pad192_pack_f9b2_v1", "arg_keys": ["x", "centroids", "x_pad", "c_pad", "B", "N", "D", "K", "total_x_pad", "total_c_pad"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 256}')) +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d160_pad192_pack_f9b2_v1", "arg_keys": ["x", "centroids", "x_pad", "c_pad", "B", "N", "D", "K", "total_x_pad", "total_c_pad"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 256}')) +pack_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d160_pad192_pack_f9b2_v1", "arg_keys": ["x", "centroids", "x_pad", "c_pad", "B", "N", "D", "K", "total_x_pad", "total_c_pad"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 256}')) + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as _common_cuda_include_dirs + return _common_cuda_include_dirs() + +def _compiled_pack_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0352"}, "kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1", 0, 256]}')) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + dtype = str(inputs.get('dtype', getattr(inputs['x'], 'dtype', 'bfloat16'))).replace('torch.', '') + if dtype not in {'bfloat16', 'bf16'}: + raise ValueError(''.join(['flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1 requires bfloat16 input, got ', format(dtype, '')])) + if dim not in SUPPORTED_D: + raise ValueError(''.join(['flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1 requires D in ', format(sorted(SUPPORTED_D), ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(n_clusters, '')])) + x_pad, c_pad = _scratch_buffers(inputs) + _launch_pack(inputs, x_pad, c_pad) + d192_inputs = _route_inputs(inputs, x_pad, c_pad) + outputs = _d192.launch_for_eval(d192_inputs) + normalized = _normalize_outputs(outputs, inputs) + route_kind = _route_kind(n_points=n_points, n_clusters=n_clusters) + selected_route = ''.join(['d', format(dim, ''), '_pad192_', format(route_kind, ''), '_repeated_mma_f9b2_v1']) + normalized['selected_route'] = selected_route + normalized['route_trace'] = {'shape_key': _shape_key(inputs), 'selected_route': selected_route, 'selected_entrypoint': ''.join([format(__name__, ''), ':launch_for_eval']), 'selected_seed': ''.join(['d', format(dim, ''), '-pad192-', format(route_kind, ''), '-repeated-mma-f9b2-v1']), 'route_kind': 'specialized', 'route_source': 'shape-specific-seed', 'guard_id': ''.join(['guard_d', format(dim, ''), '_pad192_tail_repair_f9b2_v1']), 'guard_condition': ''.join(['dtype == bfloat16 and D == ', format(dim, ''), ' and N % ', format(BLOCK_N, ''), ' == 0 and K % ', format(BLOCK_K, ''), ' == 0']), 'classification': 'seed-probe', 'dispatcher_kernel_ms': None, 'residual_contract_regions': ['zero_pad_to_d192_pack'], 'reason': 'D144/D160/D176 tail-safe route pads to D192 before the tcgen05 score producer'} + return normalized + +def _scratch_buffers(inputs: dict[str, Any]) -> tuple[Any, Any]: + import torch + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + dim = int(inputs['D']) + key = stream_cache_key(inputs, int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), bsz, n_points, n_clusters, dim, FEAT_D_PAD) + cached = _SCRATCH_CACHE.get(key) + if cached is not None: + return cached + x_pad = torch.empty((bsz, n_points, FEAT_D_PAD), dtype=inputs['x'].dtype, device=inputs['x'].device) + c_pad = torch.empty((bsz, n_clusters, FEAT_D_PAD), dtype=inputs['centroids'].dtype, device=inputs['centroids'].device) + _SCRATCH_CACHE[key] = (x_pad, c_pad) + return (x_pad, c_pad) + +def _launch_pack(inputs: dict[str, Any], x_pad: Any, c_pad: Any) -> None: + from .._dispatch_runtime import CUDAKernel + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + total_x_pad = bsz * n_points * FEAT_D_PAD + total_c_pad = bsz * n_clusters * FEAT_D_PAD + work_items = max(total_x_pad, total_c_pad) + grid_x = min((work_items + PACK_THREADS - 1) // PACK_THREADS, PACK_GRID_CAP) + cubin, kernel_name, smem_bytes, threads = _compiled_pack_kernel() + args = [inputs['x'], inputs['centroids'], x_pad, c_pad, bsz, n_points, dim, n_clusters, total_x_pad, total_c_pad] + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=(grid_x, 1, 1), block=(threads, 1, 1), args=args, shared_mem=smem_bytes) + +def _route_inputs(inputs: dict[str, Any], x_pad: Any, c_pad: Any) -> dict[str, Any]: + key = stream_cache_key(inputs, int(x_pad.data_ptr()), int(c_pad.data_ptr()), int(inputs['x_sq'].data_ptr()), int(inputs['c_sq'].data_ptr()), int(inputs['out'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), FEAT_D_PAD) + cached = _ROUTE_INPUT_CACHE.get(key) + if cached is not None: + return cached + route_inputs = {'label': inputs.get('label'), 'B': int(inputs['B']), 'N': int(inputs['N']), 'D': FEAT_D_PAD, 'K': int(inputs['K']), 'dtype': inputs.get('dtype', 'bfloat16'), 'x': x_pad, 'centroids': c_pad, 'x_sq': inputs['x_sq'], 'c_sq': inputs['c_sq'], 'out': inputs['out'], 'original_D': int(inputs['D'])} + _ROUTE_INPUT_CACHE[key] = route_inputs + return route_inputs + +def _route_kind(*, n_points: int, n_clusters: int) -> str: + if _d192._use_single_tile_path(n_points=n_points, n_clusters=n_clusters): + return 'single' + return 'paired' + +def _normalize_outputs(outputs: Any, inputs: dict[str, Any]) -> dict[str, Any]: + if outputs is None: + return {'cluster_ids': inputs['out']} + if hasattr(outputs, 'shape'): + return {'cluster_ids': outputs} + if isinstance(outputs, dict): + normalized = dict(outputs) + if 'cluster_ids' not in normalized and 'out' in normalized: + normalized['cluster_ids'] = normalized['out'] + if 'cluster_ids' in normalized: + return normalized + raise TypeError("flash_kmeans_assign D160 pad192 route must return cluster_ids or write inputs['out']") + +def _shape_key(inputs: dict[str, Any]) -> str: + label = inputs.get('label') + if label: + return str(label) + return ''.join(['b', format(int(inputs['B']), ''), '_n', format(int(inputs['N']), ''), '_k', format(int(inputs['K']), ''), '_d', format(int(inputs['D']), '')]) + +def compile_and_launch_flash_kmeans_assign_pad192(B: int=1, N: int=2048, K: int=2048, D: int=160, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(16005) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 16005}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4.py new file mode 100644 index 00000000..ef018cd6 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4.py @@ -0,0 +1,96 @@ +"""Flash-KMeans exact-D224 high-K ownership-repair candidate. + +Minimum architecture: sm_100a. This candidate uses Blackwell tcgen05/TMEM +(``lm.mma``) for an exact D=224 score producer by accumulating seven D=32 +chunks into the same TMEM score tile before the argmin role consumes it. Its +64B-swizzled TMA maps match the physical 32-element feature chunks; no host +padding or packing is used on the eval path. It is not +intended for sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +BLOCK_N = 64 +BLOCK_K = 256 +CHUNK_D = 32 +FEAT_D = 224 +D_CHUNKS = FEAT_D // CHUNK_D +SCORE_CHUNK_K = 128 +CSQ_VEC = 4 +CSQ_STAGE_VEC = 4 +ROW16_LOAD_COUNT = 64 +NUM_COMPUTE_WARPS = 4 +X_TILE_BYTES = BLOCK_N * CHUNK_D * 2 +C_TILE_BYTES = BLOCK_K * CHUNK_D * 2 +CSQ_TILE_BYTES = BLOCK_K * 4 +flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 22528, "constants": [], "cta_group": 1, "threads": 192}')) + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as _common_cuda_include_dirs + return _common_cuda_include_dirs() + +def _make_tmaps(inputs: dict[str, Any]) -> tuple[Any, Any]: + cache_key = (int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), int(inputs['D'])) + cached = inputs.get('_flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4_tmaps') + if cached is not None and cached[0] == cache_key: + return (cached[1], cached[2]) + from .._dispatch_runtime import create_tensor_map_3d_64b + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + dim = int(inputs['D']) + tmap_x = create_tensor_map_3d_64b(inputs['x'].data_ptr(), bsz * n_points, BLOCK_N, dim, CHUNK_D) + tmap_c = create_tensor_map_3d_64b(inputs['centroids'].data_ptr(), bsz * n_clusters, BLOCK_K, dim, CHUNK_D) + inputs['_flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4_tmaps'] = (cache_key, tmap_x, tmap_c) + return (tmap_x, tmap_c) + +def _compiled_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0351"}, "kernel_flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4", 22528, 192]}')) +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 22528, "constants": [], "cta_group": 1, "threads": 192}')) + +def launch_for_eval(inputs: dict[str, Any]) -> Any: + from .._dispatch_runtime import CUDAKernel + from .._dispatch_runtime import pack_kernel_args + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + if dim != FEAT_D: + raise ValueError(''.join(['flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4 requires exact D=', format(FEAT_D, ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(n_clusters, '')])) + tmap_x, tmap_c = _make_tmaps(inputs) + cubin, kernel_name, smem_bytes, threads = _compiled_kernel() + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // BLOCK_K + grid = (bsz * num_n_tiles, 1, 1) + block = (threads, 1, 1) + args = pack_kernel_args(ir, x_tmap=tmap_x, c_tmap=tmap_c, x_sq=inputs['x_sq'], c_sq=inputs['c_sq'], out=inputs['out'], B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles) + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=grid, block=block, args=args, shared_mem=smem_bytes) + return {'cluster_ids': inputs['out'], 'selected_route': 'd224_tmem_abi_repair_d17c_v4', 'route_trace': {'selected_route': 'd224_tmem_abi_repair_d17c_v4', 'selected_seed': 'flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4', 'route_kind': 'shape_specific_seed', 'guard_condition': 'dtype == bfloat16 and D == 224 and N % 64 == 0 and K % 256 == 0'}} + +def compile_and_launch_flash_kmeans_assign_splitd_d224(B: int=1, N: int=1024, K: int=512, D: int=224, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(22401) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'B': B, 'N': N, 'D': D, 'K': K, 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + passed = bool(torch.equal(out, ref)) + result: dict[str, Any] = {'passed': passed} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 22401}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d288_exactd_a532_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d288_exactd_a532_v1.py new file mode 100644 index 00000000..b1eba09b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d288_exactd_a532_v1.py @@ -0,0 +1,95 @@ +"""Flash-KMeans Euclidean assignment exact-D288 split-D candidate. + +Minimum architecture: sm_100a. This candidate uses Blackwell tcgen05/TMEM +(``lm.mma``) for an exact D=288 score producer by accumulating nine D=32 +chunks into the same TMEM score tile before the argmin role consumes it. Its +64B-swizzled TMA maps match the physical 32-element feature chunks; no host +padding or packing is used on the eval path. It is not +intended for sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +BLOCK_N = 128 +BLOCK_K = 256 +CHUNK_D = 32 +FEAT_D = 288 +D_CHUNKS = FEAT_D // CHUNK_D +SCORE_CHUNK_K = 128 +CSQ_VEC = 4 +CSQ_STAGE_VEC = 4 +NUM_COMPUTE_WARPS = 4 +X_TILE_BYTES = BLOCK_N * CHUNK_D * 2 +C_TILE_BYTES = BLOCK_K * CHUNK_D * 2 +CSQ_TILE_BYTES = BLOCK_K * 4 +flash_kmeans_assign_d288_exactd_a532_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d288_exactd_a532_v1", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 26624, "constants": [], "cta_group": 1, "threads": 192}')) + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as _common_cuda_include_dirs + return _common_cuda_include_dirs() + +def _make_tmaps(inputs: dict[str, Any]) -> tuple[Any, Any]: + cache_key = (int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), int(inputs['D'])) + cached = inputs.get('_flash_kmeans_assign_d288_exactd_a532_v1_tmaps') + if cached is not None and cached[0] == cache_key: + return (cached[1], cached[2]) + from .._dispatch_runtime import create_tensor_map_3d_64b + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + dim = int(inputs['D']) + tmap_x = create_tensor_map_3d_64b(inputs['x'].data_ptr(), bsz * n_points, BLOCK_N, dim, CHUNK_D) + tmap_c = create_tensor_map_3d_64b(inputs['centroids'].data_ptr(), bsz * n_clusters, BLOCK_K, dim, CHUNK_D) + inputs['_flash_kmeans_assign_d288_exactd_a532_v1_tmaps'] = (cache_key, tmap_x, tmap_c) + return (tmap_x, tmap_c) + +def _compiled_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0350"}, "kernel_flash_kmeans_assign_d288_exactd_a532_v1", 26624, 192]}')) +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d288_exactd_a532_v1", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 26624, "constants": [], "cta_group": 1, "threads": 192}')) + +def launch_for_eval(inputs: dict[str, Any]) -> Any: + from .._dispatch_runtime import CUDAKernel + from .._dispatch_runtime import pack_kernel_args + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + if dim != FEAT_D: + raise ValueError(''.join(['flash_kmeans_assign_d288_exactd_a532_v1 requires exact D=', format(FEAT_D, ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(n_clusters, '')])) + tmap_x, tmap_c = _make_tmaps(inputs) + cubin, kernel_name, smem_bytes, threads = _compiled_kernel() + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // BLOCK_K + grid = (bsz * num_n_tiles, 1, 1) + block = (threads, 1, 1) + args = pack_kernel_args(ir, x_tmap=tmap_x, c_tmap=tmap_c, x_sq=inputs['x_sq'], c_sq=inputs['c_sq'], out=inputs['out'], B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles) + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=grid, block=block, args=args, shared_mem=smem_bytes) + return {'cluster_ids': inputs['out'], 'selected_route': 'd288_exactd_a532_v1', 'route_trace': {'selected_route': 'd288_exactd_a532_v1', 'selected_seed': 'flash_kmeans_assign_d288_exactd_a532_v1', 'route_kind': 'shape_specific_seed', 'guard_condition': 'dtype == bfloat16 and D == 288 and N % 128 == 0 and K % 256 == 0'}} + +def compile_and_launch_flash_kmeans_assign_splitd_d288(B: int=1, N: int=1024, K: int=512, D: int=256, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(25601) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'B': B, 'N': N, 'D': D, 'K': K, 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + passed = bool(torch.equal(out, ref)) + result: dict[str, Any] = {'passed': passed} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 25601}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1.py new file mode 100644 index 00000000..0544a606 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1.py @@ -0,0 +1,88 @@ +"""Flash-KMeans Euclidean assignment D288 parent/Split-K hybrid seed. + +Minimum architecture: sm_100a. This production-dispatchable seed composes the +validated exact-D288 tcgen05/TMEM parent for ``K <= 2048`` with the validated +CTA-per-(N tile, K tile) Split-K producer plus its in-timed reduction for +``K >= 4096``. It is not intended for sm_120a/sm_121a where ptxas rejects +tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from typing import Any +from . import flash_kmeans_assign_d288_exactd_a532_v1 as _parent +from . import flash_kmeans_assign_d288_splitk_cta_0438_v1 as _splitk +BLOCK_N = _parent.BLOCK_N +BLOCK_K = _parent.BLOCK_K +SUPPORTED_D = {288} +BF16_DTYPE_NAMES = _splitk.BF16_DTYPE_NAMES +ROUTE_ID = 'd288_parent_splitk_hybrid_20260629_v1' +SEED_ID = 'd288-parent-splitk-hybrid-20260629-v1' +PARENT_CHILD_ROUTE_ID = 'd288_exactd_a532_v1' +SPLITK_CHILD_ROUTE_ID = _splitk.ROUTE_ID +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d288_exactd_a532_v1", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 26624, "constants": [], "cta_group": 1, "threads": 192}')) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz, n_points, dim, n_clusters, dtype_name = _shape_fields(inputs) + _validate_supported_shape(N=n_points, D=dim, K=n_clusters, dtype=dtype_name) + if n_clusters <= 2048: + child_outputs = _parent.launch_for_eval(inputs) + child_route = PARENT_CHILD_ROUTE_ID + child_seed = _parent.__name__.rsplit('.', maxsplit=1)[-1] + selected_path = 'parent' + elif n_clusters >= 4096: + child_outputs = _splitk.launch_for_eval(inputs) + child_route = SPLITK_CHILD_ROUTE_ID + child_seed = _splitk.SEED_ID + selected_path = 'splitk' + else: + raise ValueError(''.join(['flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1 supports K <= 2048 through the parent or K >= 4096 through Split-K, got K=', format(n_clusters, '')])) + normalized = _normalize_outputs(child_outputs, inputs) + normalized['selected_route'] = ROUTE_ID + normalized['route_trace'] = _route_trace(inputs, child_route=child_route, child_seed=child_seed, selected_path=selected_path, total_tiles=bsz * (n_points // BLOCK_N)) + return normalized + +def _validate_supported_shape(*, N: int, D: int, K: int, dtype: Any) -> None: + dtype_name = str(dtype).replace('torch.', '') + if dtype_name not in BF16_DTYPE_NAMES: + raise ValueError(''.join(['flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1 requires bfloat16 input, got ', format(dtype, '')])) + if D not in SUPPORTED_D: + raise ValueError(''.join(['flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1 requires D == 288, got ', format(D, '')])) + if N % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(N, '')])) + if K % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(K, '')])) + +def _normalize_outputs(outputs: Any, inputs: dict[str, Any]) -> dict[str, Any]: + if outputs is None: + return {'cluster_ids': inputs['out']} + if hasattr(outputs, 'shape'): + return {'cluster_ids': outputs} + if isinstance(outputs, dict): + normalized = dict(outputs) + if 'cluster_ids' not in normalized and 'out' in normalized: + normalized['cluster_ids'] = normalized['out'] + if 'cluster_ids' in normalized: + return normalized + raise TypeError("D288 parent/Split-K hybrid route must return cluster_ids or write inputs['out']") + +def _route_trace(inputs: dict[str, Any], *, child_route: str, child_seed: str, selected_path: str, total_tiles: int) -> dict[str, Any]: + return {'shape_key': _shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'specialized', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_d288_parent_splitk_hybrid_20260629_v1', 'guard_condition': 'dtype == bfloat16 and D == 288 and N % 128 == 0 and K % 256 == 0; parent child when K <= 2048, Split-K child with in-timed reduction when K >= 4096', 'classification': 'route-ok', 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'child_route': child_route, 'child_seed': child_seed, 'selected_path': selected_path, 'total_tiles': total_tiles, 'reason': 'K-guarded D288 composition of exact-D parent and real CTA Split-K reduction paths'} + +def _shape_fields(inputs: dict[str, Any]) -> tuple[int, int, int, int, str]: + return (int(inputs['B']), int(inputs['N']), int(inputs['D']), int(inputs['K']), _dtype_name(inputs)) + +def _shape_key(inputs: dict[str, Any]) -> str: + label = inputs.get('label') + if label: + return str(label) + return ''.join(['b', format(int(inputs['B']), ''), '_n', format(int(inputs['N']), ''), '_k', format(int(inputs['K']), ''), '_d', format(int(inputs['D']), '')]) + +def _dtype_name(inputs: dict[str, Any]) -> str: + dtype = inputs.get('dtype') + if dtype is not None: + return str(dtype).replace('torch.', '') + x = inputs.get('x') + if x is not None and hasattr(x, 'dtype'): + return str(x.dtype).replace('torch.', '') + return 'bfloat16' diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d288_splitk_cta_0438_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d288_splitk_cta_0438_v1.py new file mode 100644 index 00000000..a884fcf2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d288_splitk_cta_0438_v1.py @@ -0,0 +1,154 @@ +"""Flash-KMeans Euclidean assignment D288 split-K candidate. + +Minimum architecture: sm_100a. This candidate uses Blackwell tcgen05/TMEM +(``lm.mma``) for exact D=288. High-K rows split the cluster +tile loop across CTAs, write one partial argmin per K tile, then reduce those +partials with a second Weave kernel. Non-splitK rows reuse the validated D288 +parent. It is not intended for sm_120a/sm_121a where ptxas rejects +tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +import os +from functools import lru_cache +from typing import Any +from . import flash_kmeans_assign_d288_exactd_a532_v1 as _direct +from .flash_kmeans_assign_stream import stream_cache_key +BLOCK_N = _direct.BLOCK_N +BLOCK_K = _direct.BLOCK_K +CHUNK_D = _direct.CHUNK_D +D_CHUNKS = _direct.D_CHUNKS +SCORE_CHUNK_K = _direct.SCORE_CHUNK_K +CSQ_VEC = _direct.CSQ_VEC +CSQ_STAGE_VEC = _direct.CSQ_STAGE_VEC +NUM_COMPUTE_WARPS = _direct.NUM_COMPUTE_WARPS +X_TILE_BYTES = _direct.X_TILE_BYTES +C_TILE_BYTES = _direct.C_TILE_BYTES +CSQ_TILE_BYTES = _direct.CSQ_TILE_BYTES +SUPPORTED_DIMS = {288} +BF16_DTYPE_NAMES = {'bfloat16', 'torch.bfloat16'} +SPLITK_GRID_CAP = 4096 +REDUCE_THREADS = 128 +SPLITK_MIN_K_TILES = 8 +ROUTE_ID = 'd288_splitk_cta_0438_v1' +SEED_ID = 'd288-splitk-cta-0438-v1' +flash_kmeans_assign_d288_splitk_cta_0438_v1_partial = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d288_splitk_cta_0438_v1_partial", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 26624, "constants": [], "cta_group": 1, "threads": 192}')) +flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce", "arg_keys": ["partial_scores", "partial_indices", "out", "B", "N", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 128}')) +partial_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d288_splitk_cta_0438_v1_partial", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 26624, "constants": [], "cta_group": 1, "threads": 192}')) +reduce_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce", "arg_keys": ["partial_scores", "partial_indices", "out", "B", "N", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 128}')) + +def _verify_export_ir() -> Any: + if os.environ.get('LOOM_FLASH_KMEANS_D288_SPLITK_CTA_0438_V1_VERIFY_KERNEL') == 'reduce': + return reduce_ir + return partial_ir +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d288_splitk_cta_0438_v1_partial", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 26624, "constants": [], "cta_group": 1, "threads": 192}')) + +def _cuda_include_dirs() -> list[str]: + return _direct._cuda_include_dirs() + +def _compiled_partial_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0348"}, "kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial", 26624, 192]}')) + +def _compiled_reduce_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0349"}, "kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce", 0, 128]}')) + +def _compile_ir(ir_obj: Any) -> tuple[bytes, str, int, int]: + from .._dispatch_runtime import generate_kernel + from .._dispatch_runtime import compile_cuda + smem_bytes = ir_obj.computed_smem_bytes + source = generate_kernel(ir_obj, validate=False, smem_bytes=smem_bytes) + cubin = compile_cuda(source, options=['--use_fast_math'], include_dirs=_cuda_include_dirs()) + return (cubin, ''.join(['kernel_', format(ir_obj.symbol, '')]), int(smem_bytes), int(ir_obj.threads)) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + _validate_shape(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters) + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // BLOCK_K + if not _use_splitk(dim=dim, num_n_tiles=num_n_tiles, k_tiles=k_tiles): + outputs = _direct.launch_for_eval(inputs) + trace = dict(outputs.get('route_trace', {})) + trace['selected_route'] = 'd288_exactd_a532_v1' + outputs['route_trace'] = trace + return outputs + _launch_splitk(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters, num_n_tiles=num_n_tiles, k_tiles=k_tiles) + total_work = bsz * num_n_tiles * k_tiles + trace = _route_trace(inputs, total_work=total_work, grid=(min(total_work, SPLITK_GRID_CAP), 1, 1)) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': trace} + +def _validate_shape(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int) -> None: + del bsz + dtype_name = _dtype_name(inputs) + if dtype_name not in BF16_DTYPE_NAMES: + raise ValueError(''.join(['flash_kmeans_assign_d288_splitk_cta_0438_v1 requires bfloat16 input, got ', format(dtype_name, '')])) + if dim not in SUPPORTED_DIMS: + raise ValueError(''.join(['flash_kmeans_assign_d288_splitk_cta_0438_v1 supports D=', format(sorted(SUPPORTED_DIMS), ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(n_clusters, '')])) + +def _use_splitk(*, dim: int, num_n_tiles: int, k_tiles: int) -> bool: + if num_n_tiles > 16: + return False + return k_tiles >= SPLITK_MIN_K_TILES + +def _launch_splitk(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int, num_n_tiles: int, k_tiles: int) -> None: + from .._dispatch_runtime import CUDAKernel + from .._dispatch_runtime import pack_kernel_args + total_work = bsz * num_n_tiles * k_tiles + partial_scores, partial_indices = _partial_buffers(inputs, total_work) + tmap_x, tmap_c = _direct._make_tmaps(inputs) + cubin, kernel_name, smem_bytes, threads = _compiled_partial_kernel() + partial_args = pack_kernel_args(flash_kmeans_assign_d288_splitk_cta_0438_v1_partial, x_tmap=tmap_x, c_tmap=tmap_c, c_sq=inputs['c_sq'], partial_scores=partial_scores, partial_indices=partial_indices, B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles) + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=(min(total_work, SPLITK_GRID_CAP), 1, 1), block=(threads, 1, 1), args=partial_args, shared_mem=smem_bytes) + cubin, kernel_name, smem_bytes, threads = _compiled_reduce_kernel() + reduce_args = pack_kernel_args(reduce_ir, partial_scores=partial_scores, partial_indices=partial_indices, out=inputs['out'], B=bsz, N=n_points, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles) + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=(min(bsz * num_n_tiles, SPLITK_GRID_CAP), 1, 1), block=(threads, 1, 1), args=reduce_args, shared_mem=smem_bytes) + +def _partial_buffers(inputs: dict[str, Any], total_work: int) -> tuple[Any, Any]: + import torch + cache_key = stream_cache_key(inputs, int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), int(inputs['D']), int(total_work)) + cache = inputs.setdefault('_flash_kmeans_assign_d288_splitk_cta_0438_v1_partials', {}) + cached = cache.get(cache_key) + if cached is not None: + return cached + scores = torch.empty((total_work, BLOCK_N), dtype=torch.float32, device=inputs['x'].device) + indices = torch.empty((total_work, BLOCK_N), dtype=torch.int32, device=inputs['x'].device) + cache[cache_key] = (scores, indices) + return (scores, indices) + +def _route_trace(inputs: dict[str, Any], *, total_work: int, grid: tuple[int, int, int]) -> dict[str, Any]: + return {'shape_key': _shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_d288_splitk_cta_0438_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'specialized', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_d288_splitk_cta_0438_v1', 'guard_condition': 'dtype == bfloat16 and D == 288 and N % 128 == 0 and K % 256 == 0 and num_n_tiles <= 16 and K_tiles >= 8', 'classification': 'route-ok', 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'total_tiles': total_work, 'launch_grid': grid, 'reason': 'D288 split-K seed exposes K tiles as CTA work then reduces partial argmins'} + +def compile_and_launch_flash_kmeans_assign_d288_splitk_cta_0438_v1(B: int=1, N: int=512, K: int=8192, D: int=288, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(51204) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 51204}}]) + return result + +def _shape_key(inputs: dict[str, Any]) -> str: + return ''.join(['B=', format(int(inputs['B']), ''), ',N=', format(int(inputs['N']), ''), ',K=', format(int(inputs['K']), ''), ',D=', format(int(inputs['D']), '')]) + +def _dtype_name(inputs: dict[str, Any]) -> str: + return str(inputs.get('dtype', inputs['x'].dtype)).replace('torch.', '') diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d352_exactd_splitk_c95c_v2.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d352_exactd_splitk_c95c_v2.py new file mode 100644 index 00000000..570da43f --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d352_exactd_splitk_c95c_v2.py @@ -0,0 +1,34 @@ +"""Exact-D352 Split-K Flash-KMeans staging seed. + +Minimum architecture: sm_100a. This D352-only seed reuses the validated +tcgen05/TMEM D32/K256 producer-reducer ABI with D=352 (eleven D32 chunks), +avoiding the incumbent D352-to-D384 packing kernels. It is unsupported on +sm_120a/sm_121a because those targets reject tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from typing import Any +from . import flash_kmeans_assign_d480_splitk_producer_reducer_d32k256_v1 as _abi +from . import flash_kmeans_assign_gap_pad_v1 as _gap +BLOCK_N = _abi.BLOCK_N +MMA_BLOCK_K = _abi.MMA_BLOCK_K +FEAT_D = 352 +MAX_POINT_TILES = _abi.MAX_POINT_TILES +ROUTE_ID = 'd352_exactd_splitk_c95c_v2' +SEED_ID = 'd352-exactd-splitk-c95c-v2' +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d480_splitk_partial_d32k256_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 22528, "constants": [], "cta_group": 1, "threads": 192}')) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz, n, d, k = (int(inputs[x]) for x in ('B', 'N', 'D', 'K')) + dtype = str(inputs.get('dtype', getattr(inputs['x'], 'dtype', 'bfloat16'))).replace('torch.', '') + if dtype not in {'bfloat16', 'bf16'} or d != FEAT_D or n % BLOCK_N or k % MMA_BLOCK_K: + raise ValueError('D352 exact seed requires BF16 D=352, N%64=0, and K%256=0') + ntiles, ktiles = (n // BLOCK_N, k // MMA_BLOCK_K) + if ntiles > MAX_POINT_TILES or ktiles < 4: + return _gap.launch_for_eval(inputs) + _abi._launch_blockn64_splitk(inputs, bsz=bsz, n_points=n, dim=d, n_clusters=k, num_n_tiles=ntiles, k_tiles=ktiles, k_slices=ktiles) + total_work = bsz * ntiles * ktiles + trace = _abi._route_trace(inputs, total_work=total_work, grid=(min(total_work, _abi.SPLITK_GRID_CAP), 1, 1), k_slices=ktiles) + trace.update({'selected_route': ROUTE_ID, 'selected_entrypoint': ''.join([format(__name__, ''), ':launch_for_eval']), 'selected_seed': SEED_ID, 'guard_id': 'guard_d352_exactd_splitk_c95c_v2', 'guard_condition': 'D == 352 and N/64 <= 32 and K/256 >= 4', 'reason': 'exact D352 D32-chunk producer removes D384 packing'}) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': trace} diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d416_exactd_splitd_a4a579d1_v2.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d416_exactd_splitd_a4a579d1_v2.py new file mode 100644 index 00000000..3e0b11c8 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d416_exactd_splitd_a4a579d1_v2.py @@ -0,0 +1,95 @@ +"""Flash-KMeans Euclidean assignment exact-D416 split-D candidate. + +Minimum architecture: sm_100a. This candidate uses Blackwell tcgen05/TMEM +(``lm.mma``) for an exact D=416 score producer by accumulating thirteen D=32 +chunks into the same TMEM score tile before the argmin role consumes it. Its +64B-swizzled TMA maps match the physical 32-element feature chunks; no host +padding or packing is used on the eval path. It is not +intended for sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +BLOCK_N = 128 +BLOCK_K = 256 +CHUNK_D = 32 +FEAT_D = 416 +D_CHUNKS = FEAT_D // CHUNK_D +SCORE_CHUNK_K = 128 +CSQ_VEC = 4 +CSQ_STAGE_VEC = 4 +NUM_COMPUTE_WARPS = 4 +X_TILE_BYTES = BLOCK_N * CHUNK_D * 2 +C_TILE_BYTES = BLOCK_K * CHUNK_D * 2 +CSQ_TILE_BYTES = BLOCK_K * 4 +flash_kmeans_assign_d416_exactd_splitd_a4a579d1_v2 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d416_exactd_splitd_a4a579d1_v2", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 26624, "constants": [], "cta_group": 1, "threads": 192}')) + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as _common_cuda_include_dirs + return _common_cuda_include_dirs() + +def _make_tmaps(inputs: dict[str, Any]) -> tuple[Any, Any]: + cache_key = (int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), int(inputs['D'])) + cached = inputs.get('_flash_kmeans_assign_d416_exactd_splitd_a4a579d1_v2_tmaps') + if cached is not None and cached[0] == cache_key: + return (cached[1], cached[2]) + from .._dispatch_runtime import create_tensor_map_3d_64b + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + dim = int(inputs['D']) + tmap_x = create_tensor_map_3d_64b(inputs['x'].data_ptr(), bsz * n_points, BLOCK_N, dim, CHUNK_D) + tmap_c = create_tensor_map_3d_64b(inputs['centroids'].data_ptr(), bsz * n_clusters, BLOCK_K, dim, CHUNK_D) + inputs['_flash_kmeans_assign_d416_exactd_splitd_a4a579d1_v2_tmaps'] = (cache_key, tmap_x, tmap_c) + return (tmap_x, tmap_c) + +def _compiled_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0347"}, "kernel_flash_kmeans_assign_d416_exactd_splitd_a4a579d1_v2", 26624, 192]}')) +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d416_exactd_splitd_a4a579d1_v2", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 26624, "constants": [], "cta_group": 1, "threads": 192}')) + +def launch_for_eval(inputs: dict[str, Any]) -> Any: + from .._dispatch_runtime import CUDAKernel + from .._dispatch_runtime import pack_kernel_args + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + if dim % CHUNK_D: + raise ValueError(''.join(['D must be divisible by CHUNK_D=', format(CHUNK_D, ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(n_clusters, '')])) + tmap_x, tmap_c = _make_tmaps(inputs) + cubin, kernel_name, smem_bytes, threads = _compiled_kernel() + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // BLOCK_K + grid = (bsz * num_n_tiles, 1, 1) + block = (threads, 1, 1) + args = pack_kernel_args(ir, x_tmap=tmap_x, c_tmap=tmap_c, x_sq=inputs['x_sq'], c_sq=inputs['c_sq'], out=inputs['out'], B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles) + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=grid, block=block, args=args, shared_mem=smem_bytes) + return {'cluster_ids': inputs['out'], 'selected_route': 'd416_exactd_splitd_a4a579d1_v2', 'route_trace': {'selected_route': 'd416_exactd_splitd_a4a579d1_v2', 'selected_seed': 'flash_kmeans_assign_d416_exactd_splitd_a4a579d1_v2', 'route_kind': 'shape_specific_seed', 'guard_condition': 'dtype == bfloat16 and D == 416 and N % 128 == 0 and K % 256 == 0'}} + +def compile_and_launch_flash_kmeans_assign_splitd_d416(B: int=1, N: int=1024, K: int=512, D: int=416, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(41601) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'B': B, 'N': N, 'D': D, 'K': K, 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + passed = bool(torch.equal(out, ref)) + result: dict[str, Any] = {'passed': passed} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 41601}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d480_exactd_splitd_d480_exact_seed_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d480_exactd_splitd_d480_exact_seed_v1.py new file mode 100644 index 00000000..400aa5c8 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d480_exactd_splitd_d480_exact_seed_v1.py @@ -0,0 +1,30 @@ +"""Flash-KMeans exact-D480 staged seed. + +Minimum architecture: sm_100a. The contract-visible score producer is the +32-wide tcgen05/TMEM split-D core and accumulates all fifteen D=32 chunks +before the TMEM-to-register argmax writes ``cluster_ids``. This is B200/B300 +only; sm_120a/sm_121a reject tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from typing import Any +from . import flash_kmeans_assign_d416_exactd_splitd_a4a579d1_v2 as _core +BLOCK_N = _core.BLOCK_N +BLOCK_K = _core.BLOCK_K +CHUNK_D = _core.CHUNK_D +FEAT_D = 480 +ROUTE_ID = 'd480_exactd_splitd_d480_exact_seed_v1' +SEED_ID = 'd480-exactd-splitd-d480-exact-seed-v1' +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d416_exactd_splitd_a4a579d1_v2", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 26624, "constants": [], "cta_group": 1, "threads": 192}')) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + dim = int(inputs['D']) + if dim != FEAT_D: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires exact D=', format(FEAT_D, ''), ', got ', format(dim, '')])) + if int(inputs['N']) % BLOCK_N: + raise ValueError(''.join(['N must be divisible by ', format(BLOCK_N, '')])) + if int(inputs['K']) % BLOCK_K: + raise ValueError(''.join(['K must be divisible by ', format(BLOCK_K, '')])) + outputs = _core.launch_for_eval(inputs) + return {'cluster_ids': outputs['cluster_ids'], 'selected_route': ROUTE_ID, 'route_trace': {'selected_route': ROUTE_ID, 'selected_seed': SEED_ID, 'route_kind': 'shape_specific_seed', 'guard_condition': 'dtype == bfloat16 and D == 480 and N % 128 == 0 and K % 256 == 0', 'producer': 'tcgen05_mma D=32 split-D -> TMEM -> register argmax -> cluster_ids'}} diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d480_splitk_k1024_eac2_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d480_splitk_k1024_eac2_v1.py new file mode 100644 index 00000000..a5508ad8 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d480_splitk_k1024_eac2_v1.py @@ -0,0 +1,40 @@ +"""D480 K1024 Split-K activation experiment for Flash-KMeans assignment. + +Minimum architecture: sm_100a. This D480-only staging seed reuses the +validated tcgen05/TMEM D32x15 producer and Weave reducer, but exposes K=1024 +rows with at most 32 N=64 tiles to the producer/reducer ABI. The parent keeps +those rows on serial CTA ownership; this variant tests whether four K=256 +partials amortize the reducer without changing dispatcher policy. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from typing import Any +from . import flash_kmeans_assign_d480_splitk_producer_reducer_d32k256_v1 as _parent +BLOCK_N = _parent.BLOCK_N +MMA_BLOCK_K = _parent.MMA_BLOCK_K +FEAT_D = 480 +ROUTE_ID = 'd480_splitk_k1024_eac2_v1' +SEED_ID = 'd480-splitk-k1024-eac2-v1' +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d480_splitk_partial_d32k256_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 22528, "constants": [], "cta_group": 1, "threads": 192}')) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + """Run Split-K at K1024+ when the parent producer grid remains bounded.""" + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + _parent._validate_shape(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters) + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // MMA_BLOCK_K + if not _use_k1024_splitk(dim=dim, num_n_tiles=num_n_tiles, k_tiles=k_tiles): + return _parent.launch_for_eval(inputs) + k_slices = k_tiles // _parent.SPLITK_GROUP_K_TILES + _parent._launch_blockn64_splitk(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters, num_n_tiles=num_n_tiles, k_tiles=k_tiles, k_slices=k_slices) + total_work = bsz * num_n_tiles * k_slices + trace = _parent._route_trace(inputs, total_work=total_work, grid=(min(total_work, _parent.SPLITK_GRID_CAP), 1, 1), k_slices=k_slices) + trace.update({'selected_route': ROUTE_ID, 'selected_entrypoint': ''.join([format(__name__, ''), ':launch_for_eval']), 'selected_seed': SEED_ID, 'guard_id': 'guard_d480_splitk_k1024_eac2_v1', 'guard_condition': 'dtype == bfloat16 and D == 480 and N % 64 == 0 and K % 256 == 0 and N/64 <= 32 and K/256 >= 4', 'reason': 'D480 K1024+ bounded-point-tile rows use the tcgen05 partial/reducer path'}) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': trace} + +def _use_k1024_splitk(*, dim: int, num_n_tiles: int, k_tiles: int) -> bool: + return dim == FEAT_D and num_n_tiles <= _parent.MAX_POINT_TILES and (k_tiles >= 4) diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d480_splitk_producer_reducer_d32k256_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d480_splitk_producer_reducer_d32k256_v1.py new file mode 100644 index 00000000..ce8cc985 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d480_splitk_producer_reducer_d32k256_v1.py @@ -0,0 +1,172 @@ +"""Flash-KMeans exact-D480 D32/K256 Split-K producer/reducer candidate. + +Minimum architecture: sm_100a. This D480-only candidate preserves the +tcgen05/TMEM producer and exposes each K=256 centroid tile as a coarse CTA +partial; a second Weave reducer consumes the global partial-score/index ABI +and writes contract-visible ``cluster_ids``. D=480 is fifteen exact D32 +chunks, so no feature padding or fallback is used. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +import os +from functools import lru_cache +from typing import Any +from . import flash_kmeans_assign_d480_exactd_splitd_d480_exact_seed_v1 as _direct +from . import flash_kmeans_assign_highd_splitk_8de8_v1 as _base +from .flash_kmeans_assign_stream import stream_cache_key +BLOCK_N = 64 +MMA_BLOCK_K = _base.BLOCK_K +FEAT_CHUNK = 32 +ROW16_LOAD_COUNT = 64 +ROW16_LOGICAL_CHUNK_K = 128 +CSQ_STAGE_VEC = _base.CSQ_STAGE_VEC +X_TILE_BYTES = BLOCK_N * FEAT_CHUNK * 2 +C_TILE_BYTES = MMA_BLOCK_K * FEAT_CHUNK * 2 +CSQ_TILE_BYTES = MMA_BLOCK_K * 4 +SUPPORTED_DIMS = {480} +BF16_DTYPE_NAMES = set(_base.BF16_DTYPE_NAMES) +SPLITK_GRID_CAP = _base.SPLITK_GRID_CAP +SPLITK_MIN_K_TILES = _base.SPLITK_MIN_K_TILES +SPLITK_GROUP_K_TILES = 1 +SPLITK_TILE_K = MMA_BLOCK_K * SPLITK_GROUP_K_TILES +MAX_POINT_TILES = 32 +NUM_COMPUTE_WARPS = 4 +REDUCE_LANES_PER_ROW = 4 +REDUCE_THREADS = BLOCK_N * REDUCE_LANES_PER_ROW +ROUTE_ID = 'd480_splitk_producer_reducer_d32k256_v1' +SEED_ID = 'd480-splitk-producer-reducer-d32k256-v1' +VERIFY_ENV = 'LOOM_FLASH_KMEANS_D480_SPLITK_PRODUCER_REDUCER_D32K256_VERIFY_KERNEL' +flash_kmeans_assign_d480_splitk_partial_d32k256_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d480_splitk_partial_d32k256_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 22528, "constants": [], "cta_group": 1, "threads": 192}')) +flash_kmeans_assign_d480_splitk_reduce_d32k256_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d480_splitk_reduce_d32k256_v1", "arg_keys": ["partial_scores", "partial_indices", "out", "B", "N", "K", "num_n_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 256}')) +partial_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d480_splitk_partial_d32k256_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 22528, "constants": [], "cta_group": 1, "threads": 192}')) +reduce_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d480_splitk_reduce_d32k256_v1", "arg_keys": ["partial_scores", "partial_indices", "out", "B", "N", "K", "num_n_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 256}')) + +def _verify_export_ir() -> Any: + if os.environ.get(VERIFY_ENV) == 'reduce': + return reduce_ir + return partial_ir +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d480_splitk_partial_d32k256_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 22528, "constants": [], "cta_group": 1, "threads": 192}')) + +def _compiled_partial_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0345"}, "kernel_flash_kmeans_assign_d480_splitk_partial_d32k256_v1", 22528, 192]}')) + +def _compiled_reduce_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0346"}, "kernel_flash_kmeans_assign_d480_splitk_reduce_d32k256_v1", 0, 256]}')) + +@lru_cache(maxsize=1) +def _loaded_partial_kernel() -> tuple[Any, int, int]: + from .._dispatch_runtime import CUDAKernel + cubin, kernel_name, smem_bytes, threads = _compiled_partial_kernel() + return (CUDAKernel(cubin, kernel_name), smem_bytes, threads) + +@lru_cache(maxsize=1) +def _loaded_reduce_kernel() -> tuple[Any, int, int]: + from .._dispatch_runtime import CUDAKernel + cubin, kernel_name, smem_bytes, threads = _compiled_reduce_kernel() + return (CUDAKernel(cubin, kernel_name), smem_bytes, threads) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + _validate_shape(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters) + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // MMA_BLOCK_K + if not _use_blockn64_splitk(dim=dim, num_n_tiles=num_n_tiles, k_tiles=k_tiles): + outputs = _direct.launch_for_eval(inputs) + trace = dict(outputs.get('route_trace', {})) + trace['selected_route'] = _base.ROUTE_ID + trace['fallback_from'] = ROUTE_ID + trace['fallback_reason'] = 'shape does not satisfy BLOCK_N=64 grouped Split-K constraints' + outputs['route_trace'] = trace + return outputs + k_slices = k_tiles // SPLITK_GROUP_K_TILES + _launch_blockn64_splitk(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters, num_n_tiles=num_n_tiles, k_tiles=k_tiles, k_slices=k_slices) + total_work = bsz * num_n_tiles * k_slices + trace = _route_trace(inputs, total_work=total_work, grid=(min(total_work, SPLITK_GRID_CAP), 1, 1), k_slices=k_slices) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': trace} + +def _validate_shape(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int) -> None: + del bsz + dtype_name = _base._dtype_name(inputs) + if dtype_name not in BF16_DTYPE_NAMES: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires bfloat16 input, got ', format(dtype_name, '')])) + if dim not in SUPPORTED_DIMS: + raise ValueError(''.join([format(ROUTE_ID, ''), ' supports D=', format(sorted(SUPPORTED_DIMS), ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % MMA_BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by MMA_BLOCK_K=', format(MMA_BLOCK_K, ''), ', got ', format(n_clusters, '')])) + +def _use_blockn64_splitk(*, dim: int, num_n_tiles: int, k_tiles: int) -> bool: + if num_n_tiles > MAX_POINT_TILES: + return False + if k_tiles % SPLITK_GROUP_K_TILES != 0: + return False + return k_tiles >= 32 or (k_tiles >= SPLITK_MIN_K_TILES and dim >= 448) + +def _launch_blockn64_splitk(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int, num_n_tiles: int, k_tiles: int, k_slices: int) -> None: + import torch + from .._dispatch_runtime import pack_kernel_args + total_work = bsz * num_n_tiles * k_slices + partial_scores, partial_indices = _partial_buffers(inputs, total_work) + tmap_x, tmap_c = _make_tmaps(inputs) + stream = torch.cuda.current_stream() + partial_kernel, smem_bytes, threads = _loaded_partial_kernel() + partial_args = pack_kernel_args(flash_kmeans_assign_d480_splitk_partial_d32k256_v1, x_tmap=tmap_x, c_tmap=tmap_c, c_sq=inputs['c_sq'], partial_scores=partial_scores, partial_indices=partial_indices, B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles, K_slices=k_slices) + partial_kernel.launch(grid=(min(total_work, SPLITK_GRID_CAP), 1, 1), block=(threads, 1, 1), args=partial_args, shared_mem=smem_bytes, stream=stream) + reduce_kernel, smem_bytes, threads = _loaded_reduce_kernel() + reduce_args = pack_kernel_args(reduce_ir, partial_scores=partial_scores, partial_indices=partial_indices, out=inputs['out'], B=bsz, N=n_points, K=n_clusters, num_n_tiles=num_n_tiles, K_slices=k_slices) + reduce_kernel.launch(grid=(min(bsz * num_n_tiles, SPLITK_GRID_CAP), 1, 1), block=(threads, 1, 1), args=reduce_args, shared_mem=smem_bytes, stream=stream) + +def _make_tmaps(inputs: dict[str, Any]) -> tuple[Any, Any]: + cache_key = (int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), int(inputs['D']), int(BLOCK_N)) + cached = inputs.get('_flash_kmeans_assign_d480_splitk_producer_reducer_d32k256_v1_tmaps') + if cached is not None and cached[0] == cache_key: + return (cached[1], cached[2]) + from .._dispatch_runtime import create_tensor_map_3d_64b + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + dim = int(inputs['D']) + tmap_x = create_tensor_map_3d_64b(inputs['x'].data_ptr(), bsz * n_points, BLOCK_N, dim, FEAT_CHUNK) + tmap_c = create_tensor_map_3d_64b(inputs['centroids'].data_ptr(), bsz * n_clusters, MMA_BLOCK_K, dim, FEAT_CHUNK) + inputs['_flash_kmeans_assign_d480_splitk_producer_reducer_d32k256_v1_tmaps'] = (cache_key, tmap_x, tmap_c) + return (tmap_x, tmap_c) + +def _partial_buffers(inputs: dict[str, Any], total_work: int) -> tuple[Any, Any]: + import torch + cache_key = stream_cache_key(inputs, int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), int(inputs['D']), int(total_work), int(BLOCK_N), int(SPLITK_GROUP_K_TILES)) + cache = inputs.setdefault('_flash_kmeans_assign_d480_splitk_producer_reducer_d32k256_v1_partials', {}) + cached = cache.get(cache_key) + if cached is not None: + return cached + scores = torch.empty((total_work, BLOCK_N), dtype=torch.float32, device=inputs['x'].device) + indices = torch.empty((total_work, BLOCK_N), dtype=torch.int32, device=inputs['x'].device) + cache[cache_key] = (scores, indices) + return (scores, indices) + +def _route_trace(inputs: dict[str, Any], *, total_work: int, grid: tuple[int, int, int], k_slices: int) -> dict[str, Any]: + return {'shape_key': _base._shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_d480_splitk_producer_reducer_streamdep_r63_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'specialized', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_d480_splitk_producer_reducer_d32k256_v1', 'guard_condition': 'dtype == bfloat16 and D == 480 and N % 64 == 0 and K % 256 == 0 and num_n_tiles <= 32 and (K_tiles >= 32 or (K_tiles >= 16 and D >= 448))', 'classification': 'route-ok', 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'total_tiles': total_work, 'launch_grid': grid, 'tile_shape': {'BLOCK_N': BLOCK_N, 'BLOCK_K': SPLITK_TILE_K, 'mma_block_k': MMA_BLOCK_K, 'split_k_slices': k_slices, 'grouped_mma_k_tiles': SPLITK_GROUP_K_TILES, 'cluster_grouping': 1, 'tmem_layout': 'ROW_16x256B', 'producer_to_reducer_sync': 'same_stream_ordering', 'cuda_module_cache': 'persistent_per_process', 'producer_work_feed': '2x_g2_for_same_N_K'}, 'reason': 'BLOCK_N=64 G1/R4 tile-search variant uses ROW_16x256B score drains, single-MMA producer slices, and four-lane Split-K reducer'} + +def compile_and_launch_flash_kmeans_assign_d480_splitk_producer_reducer(B: int=1, N: int=512, K: int=8192, D: int=512, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(51204) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 51204}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1.py new file mode 100644 index 00000000..808dfc1a --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1.py @@ -0,0 +1,117 @@ +"""Flash-KMeans Euclidean assignment direct D64 seed. + +Minimum architecture: sm_100a. This candidate uses Blackwell tcgen05/TMEM +(``lm.mma``) for exact D=64 rows. It reuses the single 64-wide feature tile +path from the micro-D seed but TMA-loads the original tensors directly instead +of launching a BF16 pad/copy sidecar. It is not intended for sm_120a/sm_121a +where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +BLOCK_N = 128 +BLOCK_K = 256 +SCORE_CHUNK_K = 128 +CSQ_VEC = 4 +CSQ_STAGE_VEC = 4 +FEAT_D_PAD = 64 +SUPPORTED_D = {64} +NUM_COMPUTE_WARPS = 4 +X_TILE_BYTES = BLOCK_N * FEAT_D_PAD * 2 +C_TILE_BYTES = BLOCK_K * FEAT_D_PAD * 2 +CSQ_TILE_BYTES = BLOCK_K * 4 +_TMAP_CACHE: dict[tuple[int, int, int, int, int, int], Any] = {} +ROUTE_ID = 'd64_direct_single64_1p2gap_9f2a_v1' +SEED_ID = 'd64-direct-single64-1p2gap-9f2a-v1' +flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 51200, "constants": [], "cta_group": 1, "threads": 192}')) +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 51200, "constants": [], "cta_group": 1, "threads": 192}')) + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as _common_cuda_include_dirs + return _common_cuda_include_dirs() + +def _create_direct_tensor_map_3d(data_ptr: int, global_height: int, shared_height: int): + import torch + from .._dispatch_runtime import create_tensor_map_3d + device_index = torch.cuda.current_device() + key = (device_index, int(data_ptr), int(global_height), int(shared_height), FEAT_D_PAD, FEAT_D_PAD) + cached = _TMAP_CACHE.get(key) + if cached is not None: + return cached + cached = create_tensor_map_3d(data_ptr, global_height, shared_height, FEAT_D_PAD, FEAT_D_PAD).to(device=torch.device('cuda', device_index)) + _TMAP_CACHE[key] = cached + return cached + +def _make_tmaps(inputs: dict[str, Any]) -> tuple[Any, Any]: + cache_key = (int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), FEAT_D_PAD) + cached = inputs.get('_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1_tmaps') + if cached is not None and cached[0] == cache_key: + return (cached[1], cached[2]) + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + tmap_x = _create_direct_tensor_map_3d(inputs['x'].data_ptr(), bsz * n_points, BLOCK_N) + tmap_c = _create_direct_tensor_map_3d(inputs['centroids'].data_ptr(), bsz * n_clusters, BLOCK_K) + inputs['_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1_tmaps'] = (cache_key, tmap_x, tmap_c) + return (tmap_x, tmap_c) + +def _compiled_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0339"}, "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", 51200, 192]}')) + +def launch_for_eval(inputs: dict[str, Any]) -> Any: + from .._dispatch_runtime import CUDAKernel + from .._dispatch_runtime import pack_kernel_args + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + dtype = str(inputs.get('dtype', getattr(inputs['x'], 'dtype', 'bfloat16'))).replace('torch.', '') + if dtype not in {'bfloat16', 'bf16'}: + raise ValueError(''.join(['flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1 requires bfloat16 input, got ', format(dtype, '')])) + if dim not in SUPPORTED_D: + raise ValueError(''.join(['flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1 requires D=64, got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(n_clusters, '')])) + tmap_x, tmap_c = _make_tmaps(inputs) + cubin, kernel_name, smem_bytes, threads = _compiled_kernel() + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // BLOCK_K + grid = (bsz * num_n_tiles, 1, 1) + block = (threads, 1, 1) + args = pack_kernel_args(ir, x_tmap=tmap_x, c_tmap=tmap_c, x_sq=inputs['x_sq'], c_sq=inputs['c_sq'], out=inputs['out'], B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles) + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=grid, block=block, args=args, shared_mem=smem_bytes) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': _route_trace(inputs, grid=grid)} + +def compile_and_launch_flash_kmeans_assign_d64_direct(B: int=4, N: int=1024, K: int=512, D: int=64, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(6401) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 6401}}]) + return result + +def _route_trace(inputs: dict[str, Any], *, grid: tuple[int, int, int]) -> dict[str, Any]: + return {'shape_key': _shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1:launch_for_eval', 'selected_seed': SEED_ID, 'route_kind': 'specialized', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_d64_direct_single64_1p2gap_9f2a_v1', 'guard_condition': 'dtype == bfloat16 and D == 64 and N % 128 == 0 and K % 256 == 0', 'classification': 'seed-probe', 'dispatcher_kernel_ms': None, 'reason': 'direct D64 seed uses one 64-feature tcgen05 tile without BF16 pack/pad sidecar', 'launch_grid': grid} + +def _shape_key(inputs: dict[str, Any]) -> str: + label = inputs.get('label') + if label: + return str(label) + return ''.join(['b', format(int(inputs['B']), ''), '_n', format(int(inputs['N']), ''), '_k', format(int(inputs['K']), ''), '_d', format(int(inputs['D']), '')]) diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d768_no_padding_splitk_priority_d768_exact_seed_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d768_no_padding_splitk_priority_d768_exact_seed_v1.py new file mode 100644 index 00000000..a76bc69d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_d768_no_padding_splitk_priority_d768_exact_seed_v1.py @@ -0,0 +1,41 @@ +"""Exact-D768 no-padding Split-K Flash-KMeans assignment seed. + +Minimum architecture: sm_100a. This exact priority-route reuses the proven +tcgen05/TMEM BLOCK_N=64 G1/R4 producer and four-lane Weave Split-K reducer, +but extends its runtime feature-tile loop to D=768 (twelve 64-wide chunks). +The producer writes per-slice argmax state to global partial buffers; the +reducer consumes those buffers to produce contract-visible ``cluster_ids``. +No D padding or non-Weave fallback is used. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from typing import Any +from . import flash_kmeans_assign_highd_splitk_blockn64_g1r4_streamdep_r63_v1 as _g1 +TARGET_D = 768 +ROUTE_ID = 'd768_no_padding_splitk_priority_d768_exact_seed_v1' +SEED_ID = 'd768-no-padding-splitk-priority-d768-exact-seed-v1' +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_partial_blockn64_g1r4_streamdep_r63_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 192}')) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + if dim != TARGET_D: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires exact D=', format(TARGET_D, ''), ', got D=', format(dim, '')])) + dtype_name = _g1._base._dtype_name(inputs) + if dtype_name not in _g1.BF16_DTYPE_NAMES: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires bfloat16 input, got ', format(dtype_name, '')])) + if n_points % _g1.BLOCK_N != 0 or n_clusters % _g1.MMA_BLOCK_K != 0: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires N divisible by ', format(_g1.BLOCK_N, ''), ' and K divisible by ', format(_g1.MMA_BLOCK_K, ''), '; got N=', format(n_points, ''), ', K=', format(n_clusters, '')])) + num_n_tiles = n_points // _g1.BLOCK_N + k_tiles = n_clusters // _g1.MMA_BLOCK_K + if num_n_tiles > _g1.MAX_POINT_TILES or k_tiles < _g1.SPLITK_MIN_K_TILES: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires N/64 <= ', format(_g1.MAX_POINT_TILES, ''), ' and K/256 >= ', format(_g1.SPLITK_MIN_K_TILES, ''), '; got N=', format(n_points, ''), ', K=', format(n_clusters, '')])) + k_slices = k_tiles // _g1.SPLITK_GROUP_K_TILES + _g1._launch_blockn64_splitk(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters, num_n_tiles=num_n_tiles, k_tiles=k_tiles, k_slices=k_slices) + total_work = bsz * num_n_tiles * k_slices + trace = _g1._route_trace(inputs, total_work=total_work, grid=(min(total_work, _g1.SPLITK_GRID_CAP), 1, 1), k_slices=k_slices) + trace.update({'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_d768_no_padding_splitk_priority_d768_exact_seed_v1:launch_for_eval', 'selected_seed': SEED_ID, 'guard_id': 'guard_exact_d768_no_padding_splitk_priority_d768_exact_seed_v1', 'guard_condition': 'dtype == bfloat16 and D == 768 and N % 64 == 0 and K % 256 == 0', 'feature_tiles': TARGET_D // _g1.FEAT_CHUNK, 'reason': 'exact-D768 G1/R4 tcgen05 Split-K seed; twelve feature chunks with no D padding'}) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': trace} diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_dispatcher.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_dispatcher.py new file mode 100644 index 00000000..f5342678 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_dispatcher.py @@ -0,0 +1,548 @@ +"""Flash-KMeans Euclidean assignment guarded seed dispatcher. + +Minimum architecture: sm_100a. Every production route is Weave-only and uses +existing tcgen05/TMEM seed kernels; no external runtime fallback is present. +The tcgen05/TMEM routes are not intended for sm_120a/sm_121a where ptxas +rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from collections.abc import Callable +from dataclasses import dataclass +from threading import RLock +from typing import Any +from . import flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_b23d_v2 as _non128d_d160 +from . import flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_v1 as _non128d_v1 +from . import flash_kmeans_assign_cleanroom_tcgen05_v10 as _single +from . import flash_kmeans_assign_cleanroom_tcgen05_v15 as _paired +from . import flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1 as _d64_direct +from . import flash_kmeans_assign_d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1 as _d112_a262 +from . import flash_kmeans_assign_d128_splitk_priority_575c_v1 as _d128_priority +from . import flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1 as _tail_pad192_f9b2 +from . import flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4 as _d224_d17c +from . import flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1 as _d288_hybrid +from . import flash_kmeans_assign_d352_exactd_splitk_c95c_v2 as _d352_exactd +from . import flash_kmeans_assign_d480_splitk_k1024_eac2_v1 as _d480_eac2 +from . import flash_kmeans_assign_d768_no_padding_splitk_priority_d768_exact_seed_v1 as _d768_priority +from . import flash_kmeans_assign_gap_pad_v1 as _gap_pad +from . import flash_kmeans_assign_highd_paired_packedpartial_gridcap160_0194_v1 as _highd_paired_gridcap160 +from . import flash_kmeans_assign_highd_splitd_6fcf_v1 as _highd +from . import flash_kmeans_assign_highd_splitk_8de8_v1 as _highd_splitk +from . import flash_kmeans_assign_lowdim_e50c_v1 as _lowdim +from . import flash_kmeans_assign_microdim_hybrid_9c0d_v1 as _microdim +from . import flash_kmeans_assign_microdim_pipeline4_08f9_v4 as _microdim_pipeline4 +from . import flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1 as _highd_no_padding_r63 +from .flash_kmeans_assign_stream import bind_current_stream +BLOCK_N = _single.BLOCK_N +BLOCK_K = _single.BLOCK_K +FEAT_D = _single.FEAT_D +D144_D176_F9B2_DIMS = _decode_capture(_json_loads('{"__tuple__": [144, 176]}')) +SUPPORTED_DIMS = _decode_capture(_json_loads('{"__tuple__": [16, 32, 64, 48, 112, 224, 288, 352, 416, 480, 80, 96, 128, 144, 176, 160, 192, 256, 320, 384, 448, 512, 768]}')) +D64_SEED_ID = 'd64-direct-single64-1p2gap-9f2a-v1' +D112_A262_SEED_ID = _d112_a262.SEED_ID +MICRODIM_PAD64_SEED_ID = 'microdim-pad64-d64-direct-v1' +MICRODIM_HYBRID_SEED_ID = _microdim.SEED_ID +MICRODIM_PIPELINE4_SEED_ID = _microdim_pipeline4.SEED_ID +LOWDIM_E50C_SEED_ID = 'lowdim-e50c-v1' +SMALL_SEED_ID = 'small-grid-single-tile-v10' +PAIRED_SEED_ID = 'paired-large-v15' +D144_D160_D176_PAD192_F9B2_SEED_ID = 'd144-d160-d176-pad192-tail-repair-f9b2-v1' +D160_PADDED_SEED_ID = 'd160-padded_single-repeated-mma-b23d-v2' +D192_SINGLE_SEED_ID = 'd192-single-repeated-mma-v1' +D192_PAIRED_SEED_ID = 'd192-paired-repeated-mma-v1' +D256_SINGLE_SEED_ID = 'd256-single-repeated-mma-v1' +D288_PARENT_SPLITK_HYBRID_SEED_ID = _d288_hybrid.SEED_ID +HIGHD_SPLITD_SEED_ID = _highd.SEED_ID +HIGHD_SPLITK_SEED_ID = _highd_splitk.SEED_ID +HIGHD_PAIRED_PACKEDPARTIAL_SEED_ID = _highd_paired_gridcap160.SEED_ID +HIGHD_NO_PADDING_PORTFOLIO_R63_SEED_ID = _highd_no_padding_r63.SEED_ID +D768_PRIORITY_SEED_ID = _d768_priority.SEED_ID +D480_EAC2_SEED_ID = _d480_eac2.SEED_ID +GAP_PAD_SEED_ID = _gap_pad.SEED_ID +D352_EXACTD_SEED_ID = _d352_exactd.SEED_ID +D224_D17C_SEED_ID = 'flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4' +D64_ROUTE_ID = 'd64_direct_single64_1p2gap_9f2a_v1' +D112_A262_ROUTE_ID = _d112_a262.ROUTE_ID +D112_A262_LIVE_REBASE_ARTIFACT = 'artifacts/generalize_auto_tuning/generalize-auto-tuning-flash-kmeans-assign-bf8f/dispatcher_consumption/live_incumbent_rebase.json' +D112_A262_LIVE_REBASE_MEASUREMENT_SESSION_ID = 'd112-a262-live-rebase-3673c9957e9a4fba9b59c8b0a4ee076e' +D112_A262_SELECTED_SHAPE_KEYS = frozenset({(2, 2048, 112, 512), (1, 512, 112, 8192), (1, 256, 112, 256), (2, 512, 112, 8192), (1, 384, 112, 4096), (3, 768, 112, 8192)}) +D112_A262_RETAINED_SHAPE_KEYS = frozenset(_d112_a262.SHAPE_KEYS) - D112_A262_SELECTED_SHAPE_KEYS +if len(D112_A262_SELECTED_SHAPE_KEYS) != 6 or len(D112_A262_RETAINED_SHAPE_KEYS) != 8: + raise RuntimeError('D112 a262 additive ownership must preserve the measured 6/8 partition') +MICRODIM_PAD64_ROUTE_ID = 'microdim_pad64_d64_direct_v1' +MICRODIM_HYBRID_ROUTE_ID = _microdim.ROUTE_ID +MICRODIM_PIPELINE4_ROUTE_ID = _microdim_pipeline4.ROUTE_ID +LOWDIM_E50C_ROUTE_ID = 'lowdim_e50c_v1' +GAP_PAD_ROUTE_ID = _gap_pad.ROUTE_ID +D352_EXACTD_ROUTE_ID = _d352_exactd.ROUTE_ID +D224_D17C_ROUTE_ID = 'd224_tmem_abi_repair_d17c_v4' +SMALL_ROUTE_ID = 'small_grid_single_tile_v10' +PAIRED_ROUTE_ID = 'paired_large_v15' +D128_EVEN_NEAR_FLOOR_V10_ROUTE_ID = 'd128_even_near_floor_v10_repair' +D128_PRIORITY_SPLITK_ROUTE_ID = _d128_priority.ROUTE_ID +D144_D160_D176_PAD192_F9B2_ROUTE_ID = 'd144_d160_d176_pad192_tail_repair_f9b2_v1' +D160_PADDED_ROUTE_ID = 'd160_padded_single_repeated_mma_v2' +D192_SINGLE_ROUTE_ID = 'd192_single_repeated_mma_v1' +D192_PAIRED_ROUTE_ID = 'd192_paired_repeated_mma_v1' +D256_SINGLE_ROUTE_ID = 'd256_single_repeated_mma_v1' +D288_PARENT_SPLITK_HYBRID_ROUTE_ID = _d288_hybrid.ROUTE_ID +HIGHD_SPLITD_ROUTE_ID = _highd.ROUTE_ID +HIGHD_SPLITK_ROUTE_ID = _highd_splitk.ROUTE_ID +HIGHD_PAIRED_PACKEDPARTIAL_ROUTE_ID = _highd_paired_gridcap160.ROUTE_ID +HIGHD_NO_PADDING_PORTFOLIO_R63_ROUTE_ID = _highd_no_padding_r63.ROUTE_ID +D768_PRIORITY_ROUTE_ID = _d768_priority.ROUTE_ID +D480_EAC2_ROUTE_ID = _d480_eac2.ROUTE_ID +GENERIC_FALLBACK_ID = 'aligned_weave_v10_fallback' +UNSUPPORTED_ROUTE_ID = 'unsupported_shape' +SMALL_GRID_N_TILE_CAP = 8 +SMALL_GRID_K_TILE_CAP = 2 +SLOW_ROUTE_SPEEDUP_THRESHOLD = 0.98 +BF16_DTYPE_NAMES = {'bfloat16', 'bf16', 'torch.bfloat16'} +D768_PRIORITY_SHAPES = frozenset({(1, 512, 4096), (1, 1024, 4096), (1, 512, 8192), (1, 1024, 8192), (1, 2048, 4096)}) +MICRODIM_PAD64_SHAPES = frozenset({(3, 2432, 512, 16), (4, 1024, 512, 16), (3, 2432, 512, 32), (4, 1024, 512, 32)}) +MICRODIM_PIPELINE4_SOURCE_SHAPES = frozenset({(8, 65536, 16, 512), (8, 65536, 32, 512)}) +MICRODIM_PIPELINE4_SELECTED_SHAPES = MICRODIM_PIPELINE4_SOURCE_SHAPES +MICRODIM_PIPELINE4_LIVE_REBASE_ARTIFACT = 'artifacts/generalize_auto_tuning/generalize-auto-tuning-flash-kmeans-assign-a083/dispatcher_consumption/live_incumbent_rebase.json' +MICRODIM_PIPELINE4_LIVE_REBASE_MEASUREMENT_SESSION_ID = 'a083-live-rebase-80f484e6feb84482873f1cf45ed292bd' +if len(MICRODIM_PIPELINE4_SELECTED_SHAPES) != 2: + raise RuntimeError('08f9 additive ownership must preserve the measured two-row partition') +D224_D17C_SHAPES = frozenset({(4, 1536, 256, 224)}) +PRIORITY_EXPECTED_SEEDS = {128: _d128_priority.SEED_ID} +_PREPARE_LOCK = RLock() + +@dataclass(frozen=True) +class RouteDecision: + route_id: str + entrypoint: str + selected_seed: str + route_kind: str + route_source: str + guard_id: str + guard_condition: str + reason: str + + def trace_row(self, *, shape_key: str | None=None, expected_seed: str | None=None, dispatcher_kernel_ms: float | None=None, shape_specific_kernel_ms: float | None=None, relative_speedup_vs_baseline: float | None=None) -> dict[str, Any]: + return {'shape_key': shape_key, 'selected_route': self.route_id, 'selected_entrypoint': self.entrypoint, 'selected_seed': self.selected_seed, 'expected_seed': expected_seed, 'route_kind': self.route_kind, 'route_source': self.route_source, 'guard_id': self.guard_id, 'guard_condition': self.guard_condition, 'classification': _classify_route(self, expected_seed, relative_speedup_vs_baseline), 'dispatcher_kernel_ms': dispatcher_kernel_ms, 'shape_specific_kernel_ms': shape_specific_kernel_ms, 'relative_speedup_vs_baseline': relative_speedup_vs_baseline, 'reason': self.reason} + +@dataclass(frozen=True) +class PreparedKMeansLaunchPlan: + """Exact KMeans route with fully marshalled launches and owned workspace.""" + inputs: dict[str, Any] + decision: RouteDecision + arch: str + shape: tuple[int, int, int, int, str] + direct_launcher: Callable[..., Any] + launch_count: int + device_index: int + stream: Any + stream_handle: int + timeout_ms: float | None = None + + @property + def selected_route(self) -> str: + return self.decision.route_id + + @property + def launch_entrypoint(self) -> str: + return self.decision.entrypoint + + def launch(self, *, stream: Any=None, timeout_ms: float | None=None, public_inputs_already_recorded: bool=False) -> dict[str, Any]: + """Submit the bound launch sequence without dispatch or allocation.""" + import torch + if not isinstance(public_inputs_already_recorded, bool): + raise TypeError('public_inputs_already_recorded must be a bool') + with torch.cuda.device(self.device_index): + resolved_stream = self.stream if stream is None else stream + resolved_handle = int(resolved_stream.cuda_stream) + if resolved_handle != self.stream_handle: + raise RuntimeError(''.join(['prepared Flash-KMeans plan is stream-bound: prepared on stream 0x', format(self.stream_handle, ''.join(['x'])), ', requested 0x', format(resolved_handle, ''.join(['x'])), '; prepare a separate plan inside the target torch.cuda.stream(...) context'])) + try: + result = self.direct_launcher(self.inputs, stream=None, timeout_ms=self.timeout_ms if timeout_ms is None else timeout_ms) + finally: + self.direct_launcher.record_stream(resolved_stream) + if not public_inputs_already_recorded: + seen: set[int] = set() + for value in self.inputs.values(): + identity = id(value) + record_stream = getattr(value, 'record_stream', None) + if identity not in seen and callable(record_stream): + seen.add(identity) + record_stream(resolved_stream) + return result + + def _launch_stream_bound_recorded(self, *, timeout_ms: float | None=None) -> dict[str, Any]: + """Submit an internally proved fixed-stream slot without re-traversal. + + Exported stateful runtimes may call this only after validating the + prepared stream, recording every current public tensor, and recording + the plan's route-private launch storage at slot publication. + """ + return self.direct_launcher(self.inputs, stream=None, timeout_ms=self.timeout_ms if timeout_ms is None else timeout_ms) +ROUTE_SMALL_V10 = RouteDecision(route_id=SMALL_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval', selected_seed=SMALL_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_small_grid_single_tile_v10', guard_condition='dtype == bfloat16 and D == 128 and N % 128 == 0 and K % 256 == 0 and num_n_tiles <= 8 and K_tiles <= 2', reason='small-grid anchor uses the v10 single point-tile seed') +ROUTE_PAIRED_V15 = RouteDecision(route_id=PAIRED_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval', selected_seed=PAIRED_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_paired_large_v15', guard_condition='dtype == bfloat16 and D == 128 and N % 128 == 0 and K % 256 == 0 and num_n_tiles % 2 == 0 and not (num_n_tiles <= 8 and K_tiles <= 2)', reason='even point-tile grids use the v15 paired point-tile seed') +ROUTE_ALIGNED_V10_FALLBACK = RouteDecision(route_id=GENERIC_FALLBACK_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval', selected_seed=SMALL_SEED_ID, route_kind='fallback', route_source='generic-weave-fallback', guard_id='guard_aligned_v10_weave_fallback', guard_condition='dtype == bfloat16 and D == 128 and N % 128 == 0 and K % 256 == 0 and num_n_tiles % 2 == 1 and not (num_n_tiles <= 8 and K_tiles <= 2)', reason='v15 paired kernel requires an even number of point tiles; v10 is the aligned Weave fallback') +ROUTE_D128_EVEN_NEAR_FLOOR_V10_REPAIR = RouteDecision(route_id=D128_EVEN_NEAR_FLOOR_V10_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval', selected_seed=SMALL_SEED_ID, route_kind='specialized', route_source='generated-variant', guard_id='guard_d128_even_b8_n8192_k256_v10_repair', guard_condition='dtype == bfloat16 and B == 8 and D == 128 and N == 8192 and K == 256', reason='same-session post-D895 replay shows v10 beats paired v15 on the B8/N8192/K256 near-floor row') +ROUTE_D128_PRIORITY_SPLITK = RouteDecision(route_id=D128_PRIORITY_SPLITK_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_d128_splitk_priority_575c_v1:launch_for_eval', selected_seed=_d128_priority.SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_d128_no_padding_splitk_priority_575c_v1', guard_condition='B == 1 and D == 128 and N in [512,1024,2048] and K in [4096,8192]', reason='exact five-row D128 priority bucket consumes the 64x256 G1/R4 Split-K seed') +ROUTE_D64_DIRECT_9F2A = RouteDecision(route_id=D64_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1:launch_for_eval', selected_seed=D64_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_d64_direct_single64_1p2gap_9f2a_v1', guard_condition='dtype == bfloat16 and D == 64 and N % 128 == 0 and K % 256 == 0', reason='D64 lowdim-tail bucket uses the 1d9f direct one-MMA tcgen05 score producer') +ROUTE_D112_A262 = RouteDecision(route_id=D112_A262_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1:launch_for_eval', selected_seed=D112_A262_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_additive_d112_gap_pad_6_a262_bf8f_v1', guard_condition='dtype == bfloat16 and (B, N, D, K) is one of the six rows selected by the bf8f live-current-main CUPTI rebase', reason='the additive a262 portfolio owns only rows where its direct prepared GPU span cleared 0.98x of live current main; the other eight D112 rows retain the existing gap-pad route') +ROUTE_MICRODIM_PAD64 = RouteDecision(route_id=MICRODIM_PAD64_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval', selected_seed=MICRODIM_PAD64_SEED_ID, route_kind='specialized', route_source='generated-variant', guard_id='guard_microdim_pad64_d64_direct_v1', guard_condition='dtype == bfloat16 and (B,N,K,D) is one of the four retained D16/D32 K512 low-shape rows', reason='two independent same-session B200 A/B runs retained only the B3/N2432 and B4/N1024 rows for zero-padding to the existing D64 direct seed') +ROUTE_MICRODIM_HYBRID = RouteDecision(route_id=MICRODIM_HYBRID_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_microdim_hybrid_9c0d_v1:launch_for_eval', selected_seed=MICRODIM_HYBRID_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_microdim_hybrid_9c0d_v1', guard_condition='dtype == bfloat16 and D in [16, 32] and N % 128 == 0 and K % 256 == 0', reason="D16/D32 consume c92d's hybrid seed: short K=512 rows use direct staging and large/high-K rows keep 6cd2 pack+TMA") +ROUTE_MICRODIM_PIPELINE4 = RouteDecision(route_id=MICRODIM_PIPELINE4_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_microdim_pipeline4_08f9_v4:launch_for_eval', selected_seed=MICRODIM_PIPELINE4_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_additive_microdim_pipeline4_08f9_v4_a083', guard_condition='dtype == bfloat16 and (B, N, D, K) is one of the two rows selected by the a083 live-incumbent CUPTI rebase', reason='the fresh a083 three-order CUPTI ownership screen selected both exact B8/N65536/K512 D16 and D32 rows over microdim_hybrid_9c0d_v1') +ROUTE_D224_D17C = RouteDecision(route_id=D224_D17C_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4:launch_for_eval', selected_seed=D224_D17C_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_d224_d17c_exact_b4_n1536_k256', guard_condition='dtype == bfloat16 and (B,N,K,D) == (4,1536,256,224)', reason='the exact D224 single-launch seed removes the measured two-launch gap on this one validated row') +ROUTE_LOWDIM_E50C = RouteDecision(route_id=LOWDIM_E50C_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_lowdim_e50c_v1:launch_for_eval', selected_seed=LOWDIM_E50C_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_lowdim_e50c_v1_d80_d96', guard_condition='dtype == bfloat16 and D in [80, 96] and N % 128 == 0 and K % 256 == 0', reason='D80/D96 lowdim-tail bucket uses the e50c fused pad-to-128 tcgen05 route after fdac missed the same-session floor') +ROUTE_GAP_PAD = RouteDecision(route_id=GAP_PAD_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_gap_pad_v1:launch_for_eval', selected_seed=GAP_PAD_SEED_ID, route_kind='specialized', route_source='generated-variant', guard_id='guard_gap_pad_to_supported_seed_v1', guard_condition='dtype == bfloat16 and D in [48, 112, 224, 352, 416, 480] and N % 128 == 0 and K % 256 == 0', reason='between-bucket D rows zero-pad to the next supported Weave seed bucket without changing the delegated seed schedule') +ROUTE_D352_EXACTD = RouteDecision(route_id=D352_EXACTD_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_d352_exactd_splitk_c95c_v2:launch_for_eval', selected_seed=D352_EXACTD_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_d352_exactd_splitk_c95c_v2', guard_condition='dtype == bfloat16 and D == 352 and N/64 <= 32 and K/256 >= 4', reason='exact-D352 D32-chunk Split-K seed removes the D384 pack on its validated envelope') +ROUTE_D480_EAC2 = RouteDecision(route_id=D480_EAC2_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_d480_splitk_k1024_eac2_v1:launch_for_eval', selected_seed=D480_EAC2_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_d480_splitk_k1024_eac2_v1', guard_condition='dtype == bfloat16 and D == 480 and N % 64 == 0 and K % 256 == 0 and N/64 <= 32 and K/256 >= 4', reason='D480 bounded point-tile K1024+ rows consume the validated tcgen05 producer/reducer seed; low-K and large-N rows retain the Weave gap-pad route') +ROUTE_D288_PARENT_SPLITK_HYBRID = RouteDecision(route_id=D288_PARENT_SPLITK_HYBRID_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_d288_parent_splitk_hybrid_20260629_v1:launch_for_eval', selected_seed=D288_PARENT_SPLITK_HYBRID_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_d288_parent_splitk_hybrid_20260629_v1', guard_condition='dtype == bfloat16 and D == 288 and N % 128 == 0 and K % 256 == 0 and (K <= 2048 or K >= 4096)', reason='D288 consumes the measured exact-D parent/Split-K hybrid; the high-K branch includes its CTA Split-K reduction in the dispatcher-timed path') +ROUTE_D144_D176_PAD192_F9B2 = RouteDecision(route_id=D144_D160_D176_PAD192_F9B2_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_d160_pad192_tail_repair_f9b2_v1:launch_for_eval', selected_seed=D144_D160_D176_PAD192_F9B2_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_d144_d176_pad192_tail_repair_f9b2_v1', guard_condition='dtype == bfloat16 and D in [144, 176] and N % 128 == 0 and K % 256 == 0', reason='D144/D176 consume the f9b2 tail-safe seed: pack D to D=192 scratch, then reuse the D192 tcgen05 path; D160 remains on the faster promoted b23d route') +ROUTE_D160_PADDED = RouteDecision(route_id=D160_PADDED_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_b23d_v2:launch_for_eval', selected_seed=D160_PADDED_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_d160_padded_single_repeated_mma_b23d_v2', guard_condition='dtype == bfloat16 and D == 160 and N % 128 == 0 and K % 256 == 0', reason='D160 consumes the 81d5 padded-tail wrapper: pack D=160 to D=192 scratch, then reuse the D192 tcgen05 path') +ROUTE_D192_SINGLE = RouteDecision(route_id=D192_SINGLE_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_v1:launch_for_eval', selected_seed=D192_SINGLE_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_d192_single_repeated_mma_v1', guard_condition='dtype == bfloat16 and D == 192 and N % 128 == 0 and K % 256 == 0 and (num_n_tiles % 2 != 0 or (num_n_tiles <= 8 and K_tiles <= 2))', reason='D192 small or odd point-tile grids use the b23d single repeated-MMA path') +ROUTE_D192_PAIRED = RouteDecision(route_id=D192_PAIRED_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_v1:launch_for_eval', selected_seed=D192_PAIRED_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_d192_paired_repeated_mma_v1', guard_condition='dtype == bfloat16 and D == 192 and N % 128 == 0 and K % 256 == 0 and num_n_tiles % 2 == 0 and not (num_n_tiles <= 8 and K_tiles <= 2)', reason='D192 even point-tile grids use the b23d paired repeated-MMA path') +ROUTE_D256_SINGLE = RouteDecision(route_id=D256_SINGLE_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_v1:launch_for_eval', selected_seed=D256_SINGLE_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_d256_single_repeated_mma_v1', guard_condition='dtype == bfloat16 and D == 256 and N % 128 == 0 and K % 256 == 0', reason='D256 consumes the existing b23d single repeated-MMA route; paired D256 remains a resource-repair lane') +ROUTE_HIGHD_SPLITD = RouteDecision(route_id=HIGHD_SPLITD_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_highd_splitd_6fcf_v1:launch_for_eval', selected_seed=HIGHD_SPLITD_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_highd_splitd_6fcf_v1', guard_condition='dtype == bfloat16 and D in [320, 384, 448, 512] and N % 128 == 0 and K % 256 == 0 and not (num_n_tiles <= 16 and (K_tiles >= 32 or (K_tiles >= 16 and D >= 448)))', reason='high-D rows use the 6fcf split-D seed except where 8de8 split-K A/B showed wins') +ROUTE_HIGHD_SPLITK = RouteDecision(route_id=HIGHD_SPLITK_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_highd_splitk_8de8_v1:launch_for_eval', selected_seed=HIGHD_SPLITK_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_highd_splitk_8de8_v1', guard_condition='dtype == bfloat16 and D in [320, 384, 448, 512] and N % 128 == 0 and K % 256 == 0 and num_n_tiles <= 16 and (K_tiles >= 32 or (K_tiles >= 16 and D >= 448))', reason='8de8 split-K is selected only for measured low point-tile high-K high-D wins') +ROUTE_HIGHD_PAIRED_PACKEDPARTIAL = RouteDecision(route_id=HIGHD_PAIRED_PACKEDPARTIAL_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_highd_paired_packedpartial_gridcap160_0194_v1:launch_for_eval', selected_seed=HIGHD_PAIRED_PACKEDPARTIAL_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_highd_paired_packedpartial_gridcap160_0194_v1', guard_condition='dtype == bfloat16 and B == 1 and N == 2048 and K == 4096 and D in [448, 512]', reason='round-35 same-session audit showed the paired packed-partial seed beats the exported highd_splitk route for the no-padding D448/D512 paired rows') +ROUTE_HIGHD_NO_PADDING_PORTFOLIO_R63 = RouteDecision(route_id=HIGHD_NO_PADDING_PORTFOLIO_R63_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1:launch_for_eval', selected_seed=HIGHD_NO_PADDING_PORTFOLIO_R63_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_no_padding_highd_portfolio_r63_g1_d512n512_v1', guard_condition='dtype == bfloat16 and (B, N, K, D) is one of the exact no-padding high-D portfolio rows measured in R63', reason='R63 portfolio is the validated no-padding high-D path; it keeps the R52 portfolio for established rows and uses the G1/R4 streamdep child for D512/N512/K4096') +ROUTE_D768_PRIORITY = RouteDecision(route_id=D768_PRIORITY_ROUTE_ID, entrypoint='loom.examples.weave.flash_kmeans_assign_d768_no_padding_splitk_priority_d768_exact_seed_v1:launch_for_eval', selected_seed=D768_PRIORITY_SEED_ID, route_kind='specialized', route_source='shape-specific-seed', guard_id='guard_d768_no_padding_priority_exact5_incumbent', guard_condition='dtype == bfloat16 and D == 768 and (B, N, K) in [(1,512,4096), (1,1024,4096), (1,512,8192), (1,1024,8192), (1,2048,4096)]', reason='exact-D768 priority comparator route reuses the validated no-padding G1/R4 Split-K Weave seed; it closes the contract comparator gap without padding') +_ROUTE_FNS: dict[str, Callable[[dict[str, Any]], Any]] = {D64_ROUTE_ID: _d64_direct.launch_for_eval, D112_A262_ROUTE_ID: _d112_a262.launch_for_eval, MICRODIM_PAD64_ROUTE_ID: _gap_pad.launch_for_eval, MICRODIM_PIPELINE4_ROUTE_ID: _microdim_pipeline4.launch_for_eval, MICRODIM_HYBRID_ROUTE_ID: _microdim.launch_for_eval, LOWDIM_E50C_ROUTE_ID: _lowdim.launch_for_eval, D480_EAC2_ROUTE_ID: _d480_eac2.launch_for_eval, GAP_PAD_ROUTE_ID: _gap_pad.launch_for_eval, D352_EXACTD_ROUTE_ID: _d352_exactd.launch_for_eval, D224_D17C_ROUTE_ID: _d224_d17c.launch_for_eval, D288_PARENT_SPLITK_HYBRID_ROUTE_ID: _d288_hybrid.launch_for_eval, SMALL_ROUTE_ID: _single.launch_for_eval, PAIRED_ROUTE_ID: _paired.launch_for_eval, D128_PRIORITY_SPLITK_ROUTE_ID: _d128_priority.launch_for_eval, D128_EVEN_NEAR_FLOOR_V10_ROUTE_ID: _single.launch_for_eval, GENERIC_FALLBACK_ID: _single.launch_for_eval, D144_D160_D176_PAD192_F9B2_ROUTE_ID: _tail_pad192_f9b2.launch_for_eval, D160_PADDED_ROUTE_ID: _non128d_d160.launch_for_eval, D192_SINGLE_ROUTE_ID: _non128d_v1.launch_for_eval, D192_PAIRED_ROUTE_ID: _non128d_v1.launch_for_eval, D256_SINGLE_ROUTE_ID: _non128d_v1.launch_for_eval, HIGHD_SPLITD_ROUTE_ID: _highd.launch_for_eval, HIGHD_SPLITK_ROUTE_ID: _highd_splitk.launch_for_eval, HIGHD_PAIRED_PACKEDPARTIAL_ROUTE_ID: _highd_paired_gridcap160.launch_for_eval, HIGHD_NO_PADDING_PORTFOLIO_R63_ROUTE_ID: _highd_no_padding_r63.launch_for_eval, D768_PRIORITY_ROUTE_ID: _d768_priority.launch_for_eval} +_LAST_ROUTE_TRACE: list[dict[str, Any]] = [] + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + """Contract-harness entrypoint for the explicit guarded portfolio.""" + decision = select_route(inputs) + return _launch_route(inputs, decision) + +def prepare_launch_plan(inputs: dict[str, Any], *, arch: str | None=None, stream: Any=None, timeout_ms: float | None=None) -> PreparedKMeansLaunchPlan: + """Resolve and marshal a complete route once for allocation-free reuse. + + Preparation executes the selected route's host setup under launch capture. + That eagerly creates every route-specific scratch tensor, TMA descriptor, + compiled module, launch argument array, and nested child decision without + submitting GPU work. The returned plan is tied to these exact input and + output tensors. + """ + import torch + from .._dispatch_runtime import detect_gpu_arch + from .._dispatch_runtime import capture_kernel_launches + input_device = inputs['x'].device + device_index = input_device.index + if device_index is None: + device_index = torch.cuda.current_device() + device_index = int(device_index) + with torch.cuda.device(device_index): + resolved_stream = torch.cuda.current_stream(device_index) if stream is None else stream + stream_device = getattr(resolved_stream, 'device', None) + stream_device_index = getattr(stream_device, 'index', stream_device) + if stream_device_index is not None and int(stream_device_index) != device_index: + raise ValueError(''.join(['Flash-KMeans stream device ', format(stream_device_index, ''), ' does not match input device ', format(device_index, '')])) + stream_handle = int(resolved_stream.cuda_stream) + with torch.cuda.stream(resolved_stream): + detected_arch = detect_gpu_arch() + resolved_arch = detected_arch if arch is None else str(arch) + if resolved_arch != detected_arch: + raise ValueError(''.join(['flash_kmeans_assign launch arch must match the active device: requested ', format(resolved_arch, ''), ', detected ', format(detected_arch, '')])) + if resolved_arch not in {'sm_100a', 'sm_103a'}: + raise ValueError(''.join(['flash_kmeans_assign prepared routes require sm_100a or sm_103a, got ', format(resolved_arch, '')])) + decision = select_route(inputs) + shape = _shape_fields(inputs) + trace = decision.trace_row(shape_key=_shape_key(inputs)) + with _PREPARE_LOCK: + for key in tuple(inputs): + if key.startswith('_flash_kmeans_assign_'): + inputs.pop(key) + inputs['_flash_kmeans_assign_prepared_stream_key'] = (device_index, stream_handle) + inputs['_flash_kmeans_assign_dispatch_route'] = trace + route_launcher = _ROUTE_FNS[decision.route_id] + with capture_kernel_launches(stream=resolved_stream, arch=resolved_arch, inputs=inputs) as captured: + prepared_outputs = route_launcher(inputs) + prepared_result = _finish_route(inputs, decision, prepared_outputs) + direct_launcher = captured.bind(prepared_result) + _LAST_ROUTE_TRACE[:] = [trace] + return PreparedKMeansLaunchPlan(inputs=inputs, decision=decision, arch=resolved_arch, shape=shape, direct_launcher=direct_launcher, launch_count=direct_launcher.launch_count, device_index=device_index, stream=resolved_stream, stream_handle=stream_handle, timeout_ms=timeout_ms) + +def launch_prepared(plan: PreparedKMeansLaunchPlan, *, stream: Any=None, timeout_ms: float | None=None, public_inputs_already_recorded: bool=False) -> dict[str, Any]: + """Launch a plan returned by :func:`prepare_launch_plan`.""" + if not isinstance(plan, PreparedKMeansLaunchPlan): + raise TypeError('plan must be returned by prepare_launch_plan') + return plan.launch(stream=stream, timeout_ms=timeout_ms, public_inputs_already_recorded=public_inputs_already_recorded) + +def launch_for_eval_forced_fallback(inputs: dict[str, Any]) -> dict[str, Any]: + """Coverage probe entrypoint that forces the aligned Weave fallback route.""" + decision = _select_fallback_route(inputs) + return _launch_route(inputs, decision) + +def launch_for_eval_forced_d352_exactd(inputs: dict[str, Any]) -> dict[str, Any]: + """ABI preflight entrypoint that forces the validated D352 exact-D route.""" + bsz, n_points, dim, n_clusters, dtype = _shape_fields(inputs) + _validate_supported_shape(B=bsz, N=n_points, D=dim, K=n_clusters, dtype=dtype) + if not _use_d352_exactd(D=dim, num_n_tiles=n_points // BLOCK_N, k_tiles=n_clusters // BLOCK_K): + raise ValueError('forced D352 exact-D route requires D=352, N/64 <= 32, and K/256 >= 4') + return _launch_route(inputs, ROUTE_D352_EXACTD) + +def launch_for_eval_forced_d112_a262(inputs: dict[str, Any]) -> dict[str, Any]: + """ABI preflight entrypoint that forces the exact 14-row D112 portfolio.""" + bsz, n_points, dim, n_clusters, dtype = _shape_fields(inputs) + _validate_supported_shape(B=bsz, N=n_points, D=dim, K=n_clusters, dtype=dtype) + if not _supports_d112_a262(B=bsz, N=n_points, D=dim, K=n_clusters): + raise ValueError('forced D112 a262 route requires one of the exact 14 assigned D112 rows') + return _launch_route(inputs, ROUTE_D112_A262) + +def launch_for_eval_forced_microdim_pipeline4(inputs: dict[str, Any]) -> dict[str, Any]: + """ABI preflight entrypoint for both source-supported 08f9 rows.""" + bsz, n_points, dim, n_clusters, dtype = _shape_fields(inputs) + _validate_supported_shape(B=bsz, N=n_points, D=dim, K=n_clusters, dtype=dtype) + if not _supports_microdim_pipeline4(B=bsz, N=n_points, D=dim, K=n_clusters): + raise ValueError('forced microdim pipeline4 requires B=8, N=65536, K=512, and D in [16, 32]') + return _launch_route(inputs, ROUTE_MICRODIM_PIPELINE4) + +def last_route_trace() -> list[dict[str, Any]]: + return [dict(row) for row in _LAST_ROUTE_TRACE] + +def select_route(inputs: dict[str, Any]) -> RouteDecision: + bsz, n_points, dim, n_clusters, dtype = _shape_fields(inputs) + return select_route_from_shape(B=bsz, N=n_points, D=dim, K=n_clusters, dtype=dtype) + +def select_route_from_shape(*, B: int, N: int, D: int, K: int, dtype: Any='bfloat16') -> RouteDecision: + _validate_supported_shape(B=B, N=N, D=D, K=K, dtype=dtype) + num_n_tiles = N // BLOCK_N + k_tiles = K // BLOCK_K + if D == _d64_direct.FEAT_D_PAD: + return ROUTE_D64_DIRECT_9F2A + if _use_microdim_pad64(B=B, N=N, D=D, K=K): + return ROUTE_MICRODIM_PAD64 + if _use_microdim_pipeline4(B=B, N=N, D=D, K=K): + return ROUTE_MICRODIM_PIPELINE4 + if D in _microdim.SUPPORTED_D: + return ROUTE_MICRODIM_HYBRID + if D == 288 and _use_d288_parent_splitk_hybrid(K=K): + return ROUTE_D288_PARENT_SPLITK_HYBRID + if _use_d352_exactd(D=D, num_n_tiles=num_n_tiles, k_tiles=k_tiles): + return ROUTE_D352_EXACTD + if _use_d480_eac2(D=D, N=N, k_tiles=k_tiles): + return ROUTE_D480_EAC2 + if _use_d224_d17c(B=B, N=N, D=D, K=K): + return ROUTE_D224_D17C + if _use_d112_a262(B=B, N=N, D=D, K=K): + return ROUTE_D112_A262 + if D in _gap_pad.SUPPORTED_D: + return ROUTE_GAP_PAD + if D in {80, 96}: + return ROUTE_LOWDIM_E50C + if D in D144_D176_F9B2_DIMS: + return ROUTE_D144_D176_PAD192_F9B2 + if D == 160: + return ROUTE_D160_PADDED + if D == 192: + if num_n_tiles % 2 != 0 or (num_n_tiles <= SMALL_GRID_N_TILE_CAP and k_tiles <= SMALL_GRID_K_TILE_CAP): + return ROUTE_D192_SINGLE + return ROUTE_D192_PAIRED + if D == 256: + return ROUTE_D256_SINGLE + if D == _d768_priority.TARGET_D: + return ROUTE_D768_PRIORITY + if _use_d128_priority_splitk(B=B, N=N, D=D, K=K): + return ROUTE_D128_PRIORITY_SPLITK + if _use_highd_no_padding_portfolio_r63(B=B, N=N, D=D, K=K): + return ROUTE_HIGHD_NO_PADDING_PORTFOLIO_R63 + if _use_highd_paired_packedpartial(B=B, N=N, D=D, K=K): + return ROUTE_HIGHD_PAIRED_PACKEDPARTIAL + if D in _highd.SUPPORTED_DIMS: + if _use_highd_splitk(D=D, num_n_tiles=num_n_tiles, k_tiles=k_tiles): + return ROUTE_HIGHD_SPLITK + return ROUTE_HIGHD_SPLITD + if B == 8 and N == 8192 and (K == 256): + return ROUTE_D128_EVEN_NEAR_FLOOR_V10_REPAIR + if num_n_tiles <= SMALL_GRID_N_TILE_CAP and k_tiles <= SMALL_GRID_K_TILE_CAP: + return ROUTE_SMALL_V10 + if num_n_tiles % 2 == 0: + return ROUTE_PAIRED_V15 + return ROUTE_ALIGNED_V10_FALLBACK + +def route_trace_for_shape(label: str, params: dict[str, Any], *, expected_seed: str | None=None, dispatcher_kernel_ms: float | None=None, shape_specific_kernel_ms: float | None=None, relative_speedup_vs_baseline: float | None=None) -> dict[str, Any]: + try: + decision = select_route(params) + except ValueError as exc: + return _unsupported_trace_row(shape_key=label, expected_seed=expected_seed, dispatcher_kernel_ms=dispatcher_kernel_ms, shape_specific_kernel_ms=shape_specific_kernel_ms, relative_speedup_vs_baseline=relative_speedup_vs_baseline, reason=str(exc)) + return decision.trace_row(shape_key=label, expected_seed=expected_seed, dispatcher_kernel_ms=dispatcher_kernel_ms, shape_specific_kernel_ms=shape_specific_kernel_ms, relative_speedup_vs_baseline=relative_speedup_vs_baseline) + +def route_trace_for_shapes(shapes: list[dict[str, Any]], *, expected_seeds: dict[str, str] | None=None, per_shape_metrics: dict[str, dict[str, Any]] | None=None, baseline_metrics: dict[str, dict[str, Any]] | None=None) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + for shape in shapes: + label = str(shape['label']) + params = dict(shape.get('params', {})) + metrics = (per_shape_metrics or {}).get(label, {}) + baseline = (baseline_metrics or {}).get(label, {}) + dispatcher_ms = _optional_float(metrics.get('kernel_ms')) + baseline_ms = _optional_float(baseline.get('kernel_ms')) + speedup = None + if dispatcher_ms is not None and baseline_ms is not None and (dispatcher_ms > 0.0): + speedup = baseline_ms / dispatcher_ms + rows.append(route_trace_for_shape(label, params, expected_seed=(expected_seeds or {}).get(label), dispatcher_kernel_ms=dispatcher_ms, shape_specific_kernel_ms=baseline_ms, relative_speedup_vs_baseline=speedup)) + return rows + +def compile_and_launch_flash_kmeans_assign_dispatcher(B: int=2, N: int=1024, K: int=512, D: int=128, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(5101) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref)), 'route_trace': last_route_trace()} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 5101}}]) + return result + +def _launch_route(inputs: dict[str, Any], decision: RouteDecision) -> dict[str, Any]: + trace = decision.trace_row(shape_key=_shape_key(inputs)) + inputs['_flash_kmeans_assign_dispatch_route'] = trace + _LAST_ROUTE_TRACE[:] = [trace] + with bind_current_stream(inputs): + outputs = _ROUTE_FNS[decision.route_id](inputs) + return _finish_route(inputs, decision, outputs) + +def _finish_route(inputs: dict[str, Any], decision: RouteDecision, outputs: Any) -> dict[str, Any]: + trace = decision.trace_row(shape_key=_shape_key(inputs)) + normalized = _normalize_outputs(outputs, inputs) + if decision.route_id in {D112_A262_ROUTE_ID, MICRODIM_PIPELINE4_ROUTE_ID}: + child_trace = normalized.get('route_trace') + if isinstance(child_trace, dict): + trace['child_route_trace'] = dict(child_trace) + cluster_ids = normalized['cluster_ids'] + caller_out = inputs['out'] + trace['caller_owned_output'] = cluster_ids is caller_out or int(cluster_ids.data_ptr()) == int(caller_out.data_ptr()) + trace['caller_output_data_ptr'] = int(caller_out.data_ptr()) + trace['returned_output_data_ptr'] = int(cluster_ids.data_ptr()) + _LAST_ROUTE_TRACE[:] = [trace] + normalized['selected_route'] = decision.route_id + normalized['route_trace'] = trace + return normalized + +def _select_fallback_route(inputs: dict[str, Any]) -> RouteDecision: + bsz, n_points, dim, n_clusters, dtype = _shape_fields(inputs) + _validate_supported_shape(B=bsz, N=n_points, D=dim, K=n_clusters, dtype=dtype) + if dim != FEAT_D: + raise ValueError(''.join(['forced fallback is only defined for D=', format(FEAT_D, ''), ', got ', format(dim, '')])) + return ROUTE_ALIGNED_V10_FALLBACK + +def _normalize_outputs(outputs: Any, inputs: dict[str, Any]) -> dict[str, Any]: + if outputs is None: + return {'cluster_ids': inputs['out']} + if hasattr(outputs, 'shape'): + return {'cluster_ids': outputs} + if isinstance(outputs, dict): + normalized = dict(outputs) + if 'cluster_ids' not in normalized and 'out' in normalized: + normalized['cluster_ids'] = normalized['out'] + if 'cluster_ids' in normalized: + return normalized + raise TypeError("flash_kmeans_assign dispatcher route must return cluster_ids or write inputs['out']") + +def _shape_fields(inputs: dict[str, Any]) -> tuple[int, int, int, int, str]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + dtype = _dtype_name(inputs) + return (bsz, n_points, dim, n_clusters, dtype) + +def _validate_supported_shape(*, B: int, N: int, D: int, K: int, dtype: Any) -> None: + dtype_name = str(dtype).replace('torch.', '') + if dtype_name not in BF16_DTYPE_NAMES: + raise ValueError(''.join(['flash_kmeans_assign_dispatcher requires bfloat16 input, got ', format(dtype, '')])) + if D not in SUPPORTED_DIMS: + raise ValueError(''.join(['flash_kmeans_assign_dispatcher has no Weave route for D=', format(D, ''), '; supported D values are ', format(SUPPORTED_DIMS, '')])) + if D == _d768_priority.TARGET_D and (not _is_d768_priority_shape(B=B, N=N, K=K)): + raise ValueError(''.join(['flash_kmeans_assign_dispatcher has a D=768 Weave route only for the validated priority shapes ', format(sorted(D768_PRIORITY_SHAPES), ''), '; got N=', format(N, ''), ', K=', format(K, '')])) + if N % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(N, '')])) + if K % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(K, '')])) + +def _use_highd_splitk(*, D: int, num_n_tiles: int, k_tiles: int) -> bool: + if num_n_tiles > 16: + return False + return k_tiles >= 32 or (k_tiles >= _highd_splitk.SPLITK_MIN_K_TILES and D >= 448) + +def _use_d128_priority_splitk(*, B: int, N: int, D: int, K: int) -> bool: + return B == 1 and D == _d128_priority.TARGET_D and (N in (512, 1024, 2048)) and (K in (4096, 8192)) + +def _use_microdim_pad64(*, B: int, N: int, D: int, K: int) -> bool: + return (B, N, K, D) in MICRODIM_PAD64_SHAPES + +def _supports_microdim_pipeline4(*, B: int, N: int, D: int, K: int) -> bool: + return (int(B), int(N), int(D), int(K)) in MICRODIM_PIPELINE4_SOURCE_SHAPES + +def _use_microdim_pipeline4(*, B: int, N: int, D: int, K: int) -> bool: + return (int(B), int(N), int(D), int(K)) in MICRODIM_PIPELINE4_SELECTED_SHAPES + +def _use_d224_d17c(*, B: int, N: int, D: int, K: int) -> bool: + return (B, N, K, D) in D224_D17C_SHAPES + +def _supports_d112_a262(*, B: int, N: int, D: int, K: int) -> bool: + return _d112_a262.supports_shape(B=B, N=N, D=D, K=K) + +def _use_d112_a262(*, B: int, N: int, D: int, K: int) -> bool: + return (int(B), int(N), int(D), int(K)) in D112_A262_SELECTED_SHAPE_KEYS + +def _is_d768_priority_shape(*, B: int, N: int, K: int) -> bool: + return (B, N, K) in D768_PRIORITY_SHAPES + +def _use_d288_parent_splitk_hybrid(*, K: int) -> bool: + """Keep unsupported middle-K values on the existing gap-pad route.""" + return K <= 2048 or K >= 4096 + +def _use_d352_exactd(*, D: int, num_n_tiles: int, k_tiles: int) -> bool: + return D == _d352_exactd.FEAT_D and num_n_tiles <= _d352_exactd.MAX_POINT_TILES and (k_tiles >= 4) + +def _use_d480_eac2(*, D: int, N: int, k_tiles: int) -> bool: + """Apply the seed guard in eac2's 64-row point-tile units.""" + return _d480_eac2._use_k1024_splitk(dim=D, num_n_tiles=N // _d480_eac2.BLOCK_N, k_tiles=k_tiles) + +def _use_highd_paired_packedpartial(*, B: int, N: int, D: int, K: int) -> bool: + return B == 1 and N == 2048 and (K == 4096) and (D in _highd_paired_gridcap160.PAIRED_PACKEDPARTIAL_DIMS) + +def _use_highd_no_padding_portfolio_r63(*, B: int, N: int, D: int, K: int) -> bool: + return (B, N, K, D) in _highd_no_padding_r63.NO_PADDING_HIGHD_SHAPES + +def _classify_route(route: RouteDecision, expected_seed: str | None, relative_speedup_vs_baseline: float | None) -> str: + if expected_seed is not None: + if route.selected_seed != expected_seed: + return 'guard-miss' + if relative_speedup_vs_baseline is not None and relative_speedup_vs_baseline < SLOW_ROUTE_SPEEDUP_THRESHOLD: + return 'fallback-slow' if route.route_kind == 'fallback' else 'kernel-slow' + return 'seed-consumed' + if route.route_kind == 'fallback': + return 'coverage-only' + return 'route-ok' + +def _unsupported_trace_row(*, shape_key: str, expected_seed: str | None, dispatcher_kernel_ms: float | None, shape_specific_kernel_ms: float | None, relative_speedup_vs_baseline: float | None, reason: str) -> dict[str, Any]: + return {'shape_key': shape_key, 'selected_route': UNSUPPORTED_ROUTE_ID, 'selected_entrypoint': 'none', 'selected_seed': None, 'expected_seed': expected_seed, 'route_kind': 'none', 'route_source': 'unknown', 'guard_id': 'guard_miss_no_supported_weave_route', 'guard_condition': 'no production Weave guard matched this shape', 'classification': 'guard-miss' if expected_seed is not None else 'unmeasured', 'dispatcher_kernel_ms': dispatcher_kernel_ms, 'shape_specific_kernel_ms': shape_specific_kernel_ms, 'relative_speedup_vs_baseline': relative_speedup_vs_baseline, 'reason': reason} + +def _shape_key(inputs: dict[str, Any]) -> str: + label = inputs.get('label') + if label: + return str(label) + return ''.join(['b', format(int(inputs['B']), ''), '_n', format(int(inputs['N']), ''), '_k', format(int(inputs['K']), ''), '_d', format(int(inputs['D']), '')]) + +def _dtype_name(inputs: dict[str, Any]) -> str: + dtype = inputs.get('dtype') + if dtype is not None: + return str(dtype).replace('torch.', '') + x = inputs.get('x') + if x is not None and hasattr(x, 'dtype'): + return str(x.dtype).replace('torch.', '') + return 'bfloat16' + +def _optional_float(value: Any) -> float | None: + if value is None: + return None + return float(value) diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_gap_pad_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_gap_pad_v1.py new file mode 100644 index 00000000..70dfc473 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_gap_pad_v1.py @@ -0,0 +1,166 @@ +"""Flash-KMeans assignment gap-D pad-to-supported seed adapter. + +Minimum architecture: sm_100a. This generated-variant route zero-pads legal +between-bucket and selected micro-D BF16 dimensions to the next existing Weave seed bucket, then +delegates to the unmodified tcgen05/TMEM seed for that padded dimension. It is +not intended for sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +from . import flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_v1 as _non128d +from . import flash_kmeans_assign_cleanroom_tcgen05_v10 as _single +from . import flash_kmeans_assign_cleanroom_tcgen05_v15 as _paired +from . import flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1 as _d64 +from . import flash_kmeans_assign_highd_splitd_6fcf_v1 as _highd_splitd +from . import flash_kmeans_assign_highd_splitk_8de8_v1 as _highd_splitk +from .flash_kmeans_assign_stream import stream_cache_key +BLOCK_N = _single.BLOCK_N +BLOCK_K = _single.BLOCK_K +SUPPORTED_PAD: dict[int, int] = {16: 64, 32: 64, 48: 64, 112: 128, 224: 256, 288: 320, 352: 384, 416: 448, 480: 512} +SUPPORTED_D = set(SUPPORTED_PAD) +SEED_ID = 'gap-pad-to-supported-seed-v1' +ROUTE_ID = 'gap_pad_to_supported_seed_v1' +PACK_THREADS = 256 +PACK_GRID_CAP = 4096 +_SCRATCH_CACHE: dict[tuple[int, ...], tuple[Any, Any]] = {} +_ROUTE_INPUT_CACHE: dict[tuple[int, ...], dict[str, Any]] = {} +flash_kmeans_assign_gap_pad_pack_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_gap_pad_pack_v1", "arg_keys": ["x", "centroids", "x_pad", "c_pad", "B", "N", "D", "K", "D_PAD", "total_x_pad", "total_c_pad"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 256}')) +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_gap_pad_pack_v1", "arg_keys": ["x", "centroids", "x_pad", "c_pad", "B", "N", "D", "K", "D_PAD", "total_x_pad", "total_c_pad"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 256}')) +pack_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_gap_pad_pack_v1", "arg_keys": ["x", "centroids", "x_pad", "c_pad", "B", "N", "D", "K", "D_PAD", "total_x_pad", "total_c_pad"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 256}')) + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as _common_cuda_include_dirs + return _common_cuda_include_dirs() + +def _compiled_pack_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0338"}, "kernel_flash_kmeans_assign_gap_pad_pack_v1", 0, 256]}')) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + dtype = str(inputs.get('dtype', getattr(inputs['x'], 'dtype', 'bfloat16'))).replace('torch.', '') + if dtype not in {'bfloat16', 'bf16'}: + raise ValueError(''.join(['flash_kmeans_assign_gap_pad_v1 requires bfloat16 input, got ', format(dtype, '')])) + if dim not in SUPPORTED_D: + raise ValueError(''.join(['flash_kmeans_assign_gap_pad_v1 requires D in ', format(sorted(SUPPORTED_D), ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(n_clusters, '')])) + d_pad = SUPPORTED_PAD[dim] + x_pad, c_pad = _scratch_buffers(inputs, d_pad) + _launch_pack(inputs, x_pad, c_pad, d_pad) + route_inputs = _route_inputs(inputs, x_pad, c_pad, d_pad) + route_fn, route_id, route_seed, route_entrypoint = _route_for_padded_dim(d_pad=d_pad, n_points=n_points, n_clusters=n_clusters) + outputs = route_fn(route_inputs) + normalized = _normalize_outputs(outputs, inputs) + normalized['selected_route'] = ''.join(['d', format(dim, ''), '_pad', format(d_pad, ''), '_', format(route_id, '')]) + normalized['route_trace'] = {'shape_key': _shape_key(inputs), 'selected_route': normalized['selected_route'], 'selected_entrypoint': ''.join([format(__name__, ''), ':launch_for_eval']), 'selected_seed': ''.join(['d', format(dim, ''), '-pad', format(d_pad, ''), '-', format(route_seed, '')]), 'route_kind': 'specialized', 'route_source': 'generated-variant', 'guard_id': ''.join(['guard_d', format(dim, ''), '_pad', format(d_pad, ''), '_gap_pad_v1']), 'guard_condition': ''.join(['dtype == bfloat16 and D == ', format(dim, ''), ' and N % ', format(BLOCK_N, ''), ' == 0 and K % ', format(BLOCK_K, ''), ' == 0']), 'classification': 'route-ok', 'dispatcher_kernel_ms': None, 'residual_contract_regions': [''.join(['zero_pad_to_d', format(d_pad, ''), '_pack'])], 'reason': ''.join(['D', format(dim, ''), ' gap route pads to D', format(d_pad, ''), ' and delegates to ', format(route_entrypoint, '')])} + return normalized + +def _scratch_buffers(inputs: dict[str, Any], d_pad: int) -> tuple[Any, Any]: + import torch + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + dim = int(inputs['D']) + key = stream_cache_key(inputs, int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), bsz, n_points, n_clusters, dim, d_pad) + cached = _SCRATCH_CACHE.get(key) + if cached is not None: + return cached + x_pad = torch.empty((bsz, n_points, d_pad), dtype=inputs['x'].dtype, device=inputs['x'].device) + c_pad = torch.empty((bsz, n_clusters, d_pad), dtype=inputs['centroids'].dtype, device=inputs['centroids'].device) + _SCRATCH_CACHE[key] = (x_pad, c_pad) + return (x_pad, c_pad) + +def _launch_pack(inputs: dict[str, Any], x_pad: Any, c_pad: Any, d_pad: int) -> None: + from .._dispatch_runtime import CUDAKernel + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + total_x_pad = bsz * n_points * d_pad + total_c_pad = bsz * n_clusters * d_pad + work_items = max(total_x_pad, total_c_pad) + grid_x = min((work_items + PACK_THREADS - 1) // PACK_THREADS, PACK_GRID_CAP) + cubin, kernel_name, smem_bytes, threads = _compiled_pack_kernel() + args = [inputs['x'], inputs['centroids'], x_pad, c_pad, bsz, n_points, dim, n_clusters, d_pad, total_x_pad, total_c_pad] + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=(grid_x, 1, 1), block=(threads, 1, 1), args=args, shared_mem=smem_bytes) + +def _route_inputs(inputs: dict[str, Any], x_pad: Any, c_pad: Any, d_pad: int) -> dict[str, Any]: + key = stream_cache_key(inputs, int(x_pad.data_ptr()), int(c_pad.data_ptr()), int(inputs['x_sq'].data_ptr()), int(inputs['c_sq'].data_ptr()), int(inputs['out'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), d_pad) + cached = _ROUTE_INPUT_CACHE.get(key) + if cached is not None: + return cached + route_inputs = {'label': inputs.get('label'), 'B': int(inputs['B']), 'N': int(inputs['N']), 'D': d_pad, 'K': int(inputs['K']), 'dtype': inputs.get('dtype', 'bfloat16'), 'x': x_pad, 'centroids': c_pad, 'x_sq': inputs['x_sq'], 'c_sq': inputs['c_sq'], 'out': inputs['out'], 'original_D': int(inputs['D'])} + _ROUTE_INPUT_CACHE[key] = route_inputs + return route_inputs + +def _route_for_padded_dim(*, d_pad: int, n_points: int, n_clusters: int): + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // BLOCK_K + if d_pad == 64: + return (_d64.launch_for_eval, 'd64_direct_single64_1p2gap_9f2a_v1', 'd64-direct-single64-1p2gap-9f2a-v1', 'loom.examples.weave.flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1:launch_for_eval') + if d_pad == 128: + if num_n_tiles <= 8 and k_tiles <= 2: + return (_single.launch_for_eval, 'small_grid_single_tile_v10', 'small-grid-single-tile-v10', 'loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval') + if num_n_tiles % 2 == 0: + return (_paired.launch_for_eval, 'paired_large_v15', 'paired-large-v15', 'loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v15:launch_for_eval') + return (_single.launch_for_eval, 'aligned_weave_v10_fallback', 'small-grid-single-tile-v10', 'loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_v10:launch_for_eval') + if d_pad == 256: + return (_non128d.launch_for_eval, 'd256_single_repeated_mma_v1', 'd256-single-repeated-mma-v1', 'loom.examples.weave.flash_kmeans_assign_cleanroom_tcgen05_non128d_splitd_v1:launch_for_eval') + if d_pad in _highd_splitd.SUPPORTED_DIMS: + if _use_highd_splitk(d_pad=d_pad, num_n_tiles=num_n_tiles, k_tiles=k_tiles): + return (_highd_splitk.launch_for_eval, 'highd_splitk_8de8_v1', _highd_splitk.SEED_ID, 'loom.examples.weave.flash_kmeans_assign_highd_splitk_8de8_v1:launch_for_eval') + return (_highd_splitd.launch_for_eval, 'highd_splitd_6fcf_v1', _highd_splitd.SEED_ID, 'loom.examples.weave.flash_kmeans_assign_highd_splitd_6fcf_v1:launch_for_eval') + raise ValueError(''.join(['no padded route for D_PAD=', format(d_pad, '')])) + +def _use_highd_splitk(*, d_pad: int, num_n_tiles: int, k_tiles: int) -> bool: + if num_n_tiles > 16: + return False + return k_tiles >= 32 or (k_tiles >= _highd_splitk.SPLITK_MIN_K_TILES and d_pad >= 448) + +def _normalize_outputs(outputs: Any, inputs: dict[str, Any]) -> dict[str, Any]: + if outputs is None: + return {'cluster_ids': inputs['out']} + if hasattr(outputs, 'shape'): + return {'cluster_ids': outputs} + if isinstance(outputs, dict): + normalized = dict(outputs) + if 'cluster_ids' not in normalized and 'out' in normalized: + normalized['cluster_ids'] = normalized['out'] + if 'cluster_ids' in normalized: + return normalized + raise TypeError("flash_kmeans_assign gap pad route must return cluster_ids or write inputs['out']") + +def _shape_key(inputs: dict[str, Any]) -> str: + label = inputs.get('label') + if label: + return str(label) + return ''.join(['b', format(int(inputs['B']), ''), '_n', format(int(inputs['N']), ''), '_k', format(int(inputs['K']), ''), '_d', format(int(inputs['D']), '')]) + +def compile_and_launch_flash_kmeans_assign_gap_pad(B: int=2, N: int=2048, K: int=512, D: int=112, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(11205) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref)), 'route_trace': inputs.get('_flash_kmeans_assign_dispatch_route')} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 11205}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_paired_ownerreduce_r39_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_paired_ownerreduce_r39_v1.py new file mode 100644 index 00000000..93b94dda --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_paired_ownerreduce_r39_v1.py @@ -0,0 +1,158 @@ +"""Flash-KMeans high-D paired owner-reduce round-39 candidate. + +Minimum architecture: sm_100a. This additive exact-bucket seed keeps the +Blackwell tcgen05/TMEM packed-key producer on the contract-visible path for +B=1,N=2048,K=4096 high-D rows. D448 keeps the validated R1 reducer ownership +but fully unrolls its eight-slice packed-key scan after a reduce2 row-group +probe underperformed. D512 delegates to the validated gridcap160 packed-partial +route. It is not intended for sm_120a/sm_121a where ptxas rejects tcgen05 +instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +import os +from functools import lru_cache +from typing import Any +from . import flash_kmeans_assign_highd_paired_packedpartial_gridcap160_0194_v1 as _gridcap160 +from . import flash_kmeans_assign_highd_paired_packedpartial_r1_r1c1_v1 as _r1 +from . import flash_kmeans_assign_highd_paired_packedpartial_r2_7b3c_v1 as _r2 +from . import flash_kmeans_assign_highd_splitk_8de8_v1 as _base +from . import flash_kmeans_assign_highd_splitk_blockn64_g2r4_b5a6_v1 as _g2r4 +from .._dispatch_runtime import pack_kernel_args +BLOCK_N = _g2r4.BLOCK_N +MMA_BLOCK_K = _g2r4.MMA_BLOCK_K +SPLITK_GROUP_K_TILES = _g2r4.SPLITK_GROUP_K_TILES +SPLITK_TILE_K = _g2r4.SPLITK_TILE_K +SPLITK_GRID_CAP = _g2r4.SPLITK_GRID_CAP +BF16_DTYPE_NAMES = set(_base.BF16_DTYPE_NAMES) +PRODUCER_GRID_CAP_D448 = 128 +PAIRED_K_SLICES = 8 +REDUCE_ROW_GROUPS_D448 = 2 +REDUCE_ROWS_PER_CTA_D448 = BLOCK_N // REDUCE_ROW_GROUPS_D448 +REDUCE_THREADS_D448 = 32 +REDUCE_THREADS_R1_D448 = BLOCK_N +U32_MASK = 4294967295 +ROUTE_ID = 'highd_paired_ownerreduce_r39_v1' +SEED_ID = 'highd-paired-ownerreduce-r39-v1' +flash_kmeans_assign_highd_paired_ownerreduce_r39_reduce1_unroll_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_ownerreduce_r39_reduce1_unroll_v1", "arg_keys": ["partial_keys", "out", "B", "N", "K", "num_n_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 64}')) +flash_kmeans_assign_highd_paired_ownerreduce_r39_reduce2_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_ownerreduce_r39_reduce2_v1", "arg_keys": ["partial_keys", "out", "B", "N", "K", "num_n_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 32}')) +reduce1_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_ownerreduce_r39_reduce1_unroll_v1", "arg_keys": ["partial_keys", "out", "B", "N", "K", "num_n_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 64}')) +reduce2_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_ownerreduce_r39_reduce2_v1", "arg_keys": ["partial_keys", "out", "B", "N", "K", "num_n_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 32}')) +reduce_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_ownerreduce_r39_reduce1_unroll_v1", "arg_keys": ["partial_keys", "out", "B", "N", "K", "num_n_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 64}')) + +def _verify_export_ir() -> Any: + verify_kernel = os.environ.get('LOOM_FLASH_KMEANS_HIGHD_PAIRED_OWNERREDUCE_R39_VERIFY_KERNEL') + if verify_kernel == 'd448_reduce1': + return reduce1_ir + if verify_kernel == 'd448_reduce': + return reduce2_ir + if verify_kernel == 'd512_partial': + return _r2.partial_ir + if verify_kernel == 'd512_reduce': + return _r2.reduce_ir + return _r1.partial_ir +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_packedpartial_producer_r1c1_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_keys", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 192}')) + +@lru_cache(maxsize=1) +def _loaded_reduce2_kernel() -> tuple[Any, int, int]: + from .._dispatch_runtime import CUDAKernel + cubin, kernel_name, smem_bytes, threads = _base._compile_ir(reduce2_ir) + return (CUDAKernel(cubin, kernel_name), smem_bytes, threads) + +@lru_cache(maxsize=1) +def _loaded_reduce1_kernel() -> tuple[Any, int, int]: + from .._dispatch_runtime import CUDAKernel + cubin, kernel_name, smem_bytes, threads = _base._compile_ir(reduce1_ir) + return (CUDAKernel(cubin, kernel_name), smem_bytes, threads) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + _validate_shape(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters) + if _use_d448_custom(bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters): + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // MMA_BLOCK_K + k_slices = k_tiles // SPLITK_GROUP_K_TILES + if k_slices != PAIRED_K_SLICES: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires K_slices=', format(PAIRED_K_SLICES, ''), ', got ', format(k_slices, '')])) + total_work = bsz * num_n_tiles * k_slices + producer_grid, reducer_grid = _launch_d448_reduce1(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters, num_n_tiles=num_n_tiles, k_tiles=k_tiles, k_slices=k_slices, total_work=total_work) + trace = _route_trace_d448(inputs, total_work=total_work, producer_grid=producer_grid, reducer_grid=reducer_grid, k_slices=k_slices) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': trace} + outputs = _gridcap160.launch_for_eval(inputs) + trace = _wrap_gridcap160_trace(inputs, child_outputs=outputs) + return {'cluster_ids': outputs['cluster_ids'], 'selected_route': ROUTE_ID, 'route_trace': trace} + +def _use_d448_custom(*, bsz: int, n_points: int, dim: int, n_clusters: int) -> bool: + return bsz == 1 and n_points == 2048 and (n_clusters == 4096) and (dim == 448) + +def _validate_shape(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int) -> None: + dtype_name = _base._dtype_name(inputs) + if dtype_name not in BF16_DTYPE_NAMES: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires bfloat16 input, got ', format(dtype_name, '')])) + if (bsz, n_points, n_clusters, dim) not in {(1, 2048, 4096, 448), (1, 2048, 4096, 512)}: + raise ValueError(''.join([format(ROUTE_ID, ''), ' is exact-shape only for paired D448/D512 rows'])) + +def _launch_d448_reduce2(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int, num_n_tiles: int, k_tiles: int, k_slices: int, total_work: int) -> tuple[tuple[int, int, int], tuple[int, int, int]]: + import torch + partial_keys = _r1._partial_key_buffer(inputs, total_work) + tmap_x, tmap_c = _g2r4._make_tmaps(inputs) + stream = torch.cuda.current_stream() + partial_kernel, smem_bytes, threads = _r1._loaded_partial_key_kernel() + partial_args = pack_kernel_args(_r1.partial_ir, x_tmap=tmap_x, c_tmap=tmap_c, c_sq=inputs['c_sq'], partial_keys=partial_keys, B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles, K_slices=k_slices) + producer_grid = (min(total_work, PRODUCER_GRID_CAP_D448), 1, 1) + partial_kernel.launch(grid=producer_grid, block=(threads, 1, 1), args=partial_args, shared_mem=smem_bytes, stream=stream) + reduce_kernel, smem_bytes, threads = _loaded_reduce2_kernel() + reduce_args = [partial_keys, inputs['out'], bsz, n_points, n_clusters, num_n_tiles, k_slices] + reducer_work = bsz * num_n_tiles * REDUCE_ROW_GROUPS_D448 + reducer_grid = (min(reducer_work, SPLITK_GRID_CAP), 1, 1) + reduce_kernel.launch(grid=reducer_grid, block=(threads, 1, 1), args=reduce_args, shared_mem=smem_bytes, stream=stream) + return (producer_grid, reducer_grid) + +def _launch_d448_reduce1(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int, num_n_tiles: int, k_tiles: int, k_slices: int, total_work: int) -> tuple[tuple[int, int, int], tuple[int, int, int]]: + import torch + partial_keys = _r1._partial_key_buffer(inputs, total_work) + tmap_x, tmap_c = _g2r4._make_tmaps(inputs) + stream = torch.cuda.current_stream() + partial_kernel, smem_bytes, threads = _r1._loaded_partial_key_kernel() + partial_args = pack_kernel_args(_r1.partial_ir, x_tmap=tmap_x, c_tmap=tmap_c, c_sq=inputs['c_sq'], partial_keys=partial_keys, B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles, K_slices=k_slices) + producer_grid = (min(total_work, PRODUCER_GRID_CAP_D448), 1, 1) + partial_kernel.launch(grid=producer_grid, block=(threads, 1, 1), args=partial_args, shared_mem=smem_bytes, stream=stream) + reduce_kernel, smem_bytes, threads = _loaded_reduce1_kernel() + reduce_args = [partial_keys, inputs['out'], bsz, n_points, n_clusters, num_n_tiles, k_slices] + reducer_grid = (min(bsz * num_n_tiles, SPLITK_GRID_CAP), 1, 1) + reduce_kernel.launch(grid=reducer_grid, block=(threads, 1, 1), args=reduce_args, shared_mem=smem_bytes, stream=stream) + return (producer_grid, reducer_grid) + +def _route_trace_d448(inputs: dict[str, Any], *, total_work: int, producer_grid: tuple[int, int, int], reducer_grid: tuple[int, int, int], k_slices: int) -> dict[str, Any]: + return {'shape_key': _base._shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_highd_paired_ownerreduce_r39_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'shape-specific-seed', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_highd_paired_ownerreduce_r39_v1', 'guard_condition': 'dtype == bfloat16 and B == 1 and N == 2048 and K == 4096 and D in [448,512]; D448 uses unrolled R1 reducer ownership, D512 delegates to gridcap160', 'classification': 'route-ok', 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'total_tiles': total_work, 'launch_grid': producer_grid, 'reduce_grid': reducer_grid, 'paired_packedpartial_route': True, 'selected_reducer_variant': 'r39_reduce1_unroll8', 'tile_shape': {'BLOCK_N': BLOCK_N, 'BLOCK_K': SPLITK_TILE_K, 'mma_block_k': MMA_BLOCK_K, 'split_k_slices': k_slices, 'grouped_mma_k_tiles': SPLITK_GROUP_K_TILES, 'producer_work_items': total_work, 'producer_grid_cap': PRODUCER_GRID_CAP_D448, 'producer_grid_work_items_per_cta': total_work // min(total_work, PRODUCER_GRID_CAP_D448), 'producer_to_consumer': 'u64_partial_key_buffer_plus_reduce1_unroll8', 'partial_key': 'ordered_f32_score_high32_inverse_cluster_index_low32', 'reducer_lanes_per_row': 1, 'reducer_threads': REDUCE_THREADS_R1_D448, 'producer_to_reducer_sync': 'same_stream_ordering', 'reducer_scan_unroll': PAIRED_K_SLICES}, 'reason': 'D448 keeps the R1 tcgen05 packed-key producer and row ownership but fully unrolls the eight-slice packed-key scan after the reduce2 row-group probe underfed the seed'} + +def _wrap_gridcap160_trace(inputs: dict[str, Any], *, child_outputs: dict[str, Any]) -> dict[str, Any]: + child_trace = dict(child_outputs.get('route_trace', {})) + trace = dict(child_trace) + trace.update({'shape_key': _base._shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_highd_paired_ownerreduce_r39_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'shape-specific-seed-composite', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_highd_paired_ownerreduce_r39_v1', 'guard_condition': 'dtype == bfloat16 and B == 1 and N == 2048 and K == 4096 and D in [448,512]; D448 uses unrolled R1 reducer ownership, D512 delegates to gridcap160', 'classification': 'route-ok', 'child_route': child_outputs.get('selected_route') or child_trace.get('selected_route'), 'child_entrypoint': child_trace.get('selected_entrypoint'), 'child_guard_id': child_trace.get('guard_id'), 'child_tile_shape': child_trace.get('tile_shape'), 'selected_reducer_variant': 'gridcap160_child', 'paired_packedpartial_route': True, 'reason': 'D512 keeps the validated gridcap160 producer cap and R2 reducer child'}) + return trace + +def compile_and_launch_flash_kmeans_assign_highd_paired_ownerreduce_r39(B: int=1, N: int=2048, K: int=4096, D: int=448, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + seed = 44801 if D == 448 else 51202 + torch.manual_seed(seed) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': seed}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_paired_packedpartial_gridcap128_8512_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_paired_packedpartial_gridcap128_8512_v1.py new file mode 100644 index 00000000..844dc6cb --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_paired_packedpartial_gridcap128_8512_v1.py @@ -0,0 +1,134 @@ +"""Flash-KMeans high-D paired packed-partial grid-cap candidate. + +Minimum architecture: sm_100a. This additive bucket-kernel candidate keeps the +Blackwell tcgen05/TMEM G2 packed-key distance producer on the contract-visible +path for paired B=1,N=2048,K=4096 high-D rows. It preserves the validated R1 +reducer for D=448 and R2 reducer for D=512, but caps producer launch work at +128 CTAs so each producer CTA owns two fixed K-slice work items. Huge-K focus +rows delegate to the validated 4f2c stream-ordered child. It is not intended +for sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +import os +from typing import Any +from . import flash_kmeans_assign_highd_paired_packedpartial_r1_r1c1_v1 as _r1 +from . import flash_kmeans_assign_highd_paired_packedpartial_r2_7b3c_v1 as _r2 +from . import flash_kmeans_assign_highd_splitk_8de8_v1 as _base +from . import flash_kmeans_assign_highd_splitk_blockn64_g2r4_b5a6_v1 as _g2r4 +from . import flash_kmeans_assign_highd_splitk_blockn64_g2r4_streamdep_4f2c_v1 as _streamdep +from .._dispatch_runtime import pack_kernel_args +BLOCK_N = _g2r4.BLOCK_N +MMA_BLOCK_K = _g2r4.MMA_BLOCK_K +SPLITK_GROUP_K_TILES = _g2r4.SPLITK_GROUP_K_TILES +SPLITK_TILE_K = _g2r4.SPLITK_TILE_K +SPLITK_GRID_CAP = _g2r4.SPLITK_GRID_CAP +PRODUCER_GRID_CAP = 128 +SUPPORTED_DIMS = set(_base.SUPPORTED_DIMS) +BF16_DTYPE_NAMES = set(_base.BF16_DTYPE_NAMES) +PAIRED_PACKEDPARTIAL_DIMS = {448, 512} +ROUTE_ID = 'highd_paired_packedpartial_gridcap128_8512_v1' +SEED_ID = 'highd-paired-packedpartial-gridcap128-8512-v1' + +def _verify_export_ir() -> Any: + verify_kernel = os.environ.get('LOOM_FLASH_KMEANS_HIGHD_PAIRED_PACKEDPARTIAL_GRIDCAP128_8512_VERIFY_KERNEL') + if verify_kernel == 'r1_reduce': + return _r1.reduce_ir + if verify_kernel == 'r2_partial': + return _r2.partial_ir + if verify_kernel == 'r2_reduce': + return _r2.reduce_ir + if verify_kernel == 'streamdep_partial': + return _streamdep.partial_ir + if verify_kernel == 'streamdep_reduce': + return _streamdep.reduce_ir + return _r1.partial_ir +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_packedpartial_producer_r1c1_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_keys", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 192}')) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + _validate_shape(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters) + if _use_paired_gridcap(bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters): + child = _r1 if dim == 448 else _r2 + reducer = 'r1_single_lane' if dim == 448 else 'r2_two_lane' + _launch_paired_gridcap(inputs, child=child, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters) + trace = _route_trace_paired(inputs, child=child, reducer=reducer) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': trace} + outputs = _streamdep.launch_for_eval(inputs) + trace = _wrap_child_trace(inputs, child_outputs=outputs) + return {'cluster_ids': outputs['cluster_ids'], 'selected_route': ROUTE_ID, 'route_trace': trace} + +def _use_paired_gridcap(*, bsz: int, n_points: int, dim: int, n_clusters: int) -> bool: + return bsz == 1 and n_points == 2048 and (n_clusters == 4096) and (dim in PAIRED_PACKEDPARTIAL_DIMS) + +def _validate_shape(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int) -> None: + del bsz + dtype_name = _base._dtype_name(inputs) + if dtype_name not in BF16_DTYPE_NAMES: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires bfloat16 input, got ', format(dtype_name, '')])) + if dim not in SUPPORTED_DIMS: + raise ValueError(''.join([format(ROUTE_ID, ''), ' supports D=', format(sorted(SUPPORTED_DIMS), ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % MMA_BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by MMA_BLOCK_K=', format(MMA_BLOCK_K, ''), ', got ', format(n_clusters, '')])) + +def _launch_paired_gridcap(inputs: dict[str, Any], *, child: Any, bsz: int, n_points: int, dim: int, n_clusters: int) -> None: + import torch + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // MMA_BLOCK_K + k_slices = k_tiles // SPLITK_GROUP_K_TILES + total_work = bsz * num_n_tiles * k_slices + partial_keys = child._partial_key_buffer(inputs, total_work) + tmap_x, tmap_c = _g2r4._make_tmaps(inputs) + stream = torch.cuda.current_stream() + partial_kernel, smem_bytes, threads = child._loaded_partial_key_kernel() + partial_args = pack_kernel_args(child.partial_ir, x_tmap=tmap_x, c_tmap=tmap_c, c_sq=inputs['c_sq'], partial_keys=partial_keys, B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles, K_slices=k_slices) + partial_kernel.launch(grid=(min(total_work, PRODUCER_GRID_CAP), 1, 1), block=(threads, 1, 1), args=partial_args, shared_mem=smem_bytes, stream=stream) + reduce_kernel, smem_bytes, threads = child._loaded_reduce_key_kernel() + reduce_args = [partial_keys, inputs['out'], bsz, n_points, n_clusters, num_n_tiles, k_slices] + reduce_kernel.launch(grid=(min(bsz * num_n_tiles, SPLITK_GRID_CAP), 1, 1), block=(threads, 1, 1), args=reduce_args, shared_mem=smem_bytes, stream=stream) + +def _route_trace_paired(inputs: dict[str, Any], *, child: Any, reducer: str) -> dict[str, Any]: + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // MMA_BLOCK_K + k_slices = k_tiles // SPLITK_GROUP_K_TILES + total_work = int(inputs['B']) * num_n_tiles * k_slices + child_entrypoint = 'loom.examples.weave.flash_kmeans_assign_highd_paired_packedpartial_r1_r1c1_v1:launch_for_eval' if child is _r1 else 'loom.examples.weave.flash_kmeans_assign_highd_paired_packedpartial_r2_7b3c_v1:launch_for_eval' + child_route = _r1.ROUTE_ID if child is _r1 else _r2.ROUTE_ID + child_guard = 'guard_highd_paired_packedpartial_r1_r1c1_v1' if child is _r1 else 'guard_highd_paired_packedpartial_r2_7b3c_v1' + return {'shape_key': _base._shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_highd_paired_packedpartial_gridcap128_8512_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'shape-specific-seed-composite', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_highd_paired_packedpartial_gridcap128_8512_v1', 'guard_condition': 'dtype == bfloat16 and B == 1 and N == 2048 and K == 4096 and D in [448,512]; D=448 uses packed partial R1, D=512 uses packed partial R2, producer grid capped at 128 CTAs', 'classification': 'route-ok', 'child_route': child_route, 'child_entrypoint': child_entrypoint, 'child_guard_id': child_guard, 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'total_tiles': total_work, 'launch_grid': (min(total_work, PRODUCER_GRID_CAP), 1, 1), 'paired_packedpartial_route': True, 'selected_reducer_variant': reducer, 'tile_shape': {'BLOCK_N': BLOCK_N, 'BLOCK_K': SPLITK_TILE_K, 'mma_block_k': MMA_BLOCK_K, 'split_k_slices': k_slices, 'grouped_mma_k_tiles': SPLITK_GROUP_K_TILES, 'producer_work_items': total_work, 'producer_grid_cap': PRODUCER_GRID_CAP, 'producer_grid_work_items_per_cta': total_work // min(total_work, PRODUCER_GRID_CAP), 'producer_to_consumer': 'u64_partial_key_buffer_plus_shape_guarded_r1_or_r2_reduce', 'partial_key': 'ordered_f32_score_high32_inverse_cluster_index_low32', 'reducer_lanes_per_row': _r1.REDUCE_LANES_PER_ROW if child is _r1 else _r2.REDUCE_LANES_PER_ROW, 'producer_to_reducer_sync': 'same_stream_ordering'}, 'reason': 'paired high-D row keeps the G2 tcgen05 packed-key producer but caps producer CTAs at 128 to reduce launch wave scheduling overhead while preserving the same reducer ABI'} + +def _wrap_child_trace(inputs: dict[str, Any], *, child_outputs: dict[str, Any]) -> dict[str, Any]: + child_trace = dict(child_outputs.get('route_trace', {})) + child_route = child_outputs.get('selected_route') or child_trace.get('selected_route') + trace = dict(child_trace) + trace.update({'shape_key': _base._shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_highd_paired_packedpartial_gridcap128_8512_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'shape-specific-seed-composite', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_highd_paired_packedpartial_gridcap128_8512_v1', 'guard_condition': 'dtype == bfloat16 and D in [384,448,512] and N % 64 == 0 and K % 256 == 0; paired B=1,N=2048,K=4096,D in [448,512] uses 128-grid-cap packed partial keys, other high-D rows use highd_splitk_blockn64_g2r4_streamdep_4f2c_v1', 'classification': 'route-ok', 'child_route': child_route, 'child_entrypoint': child_trace.get('selected_entrypoint'), 'child_guard_id': child_trace.get('guard_id'), 'child_tile_shape': child_trace.get('tile_shape'), 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'paired_packedpartial_route': False, 'reason': 'non-paired high-D focus row keeps the stream-ordered G2/R4 Split-K child'}) + return trace + +def compile_and_launch_flash_kmeans_assign_highd_paired_packedpartial(B: int=1, N: int=2048, K: int=4096, D: int=512, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + seed = 44801 if D == 448 else 51202 + torch.manual_seed(seed) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': seed}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_paired_packedpartial_gridcap160_0194_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_paired_packedpartial_gridcap160_0194_v1.py new file mode 100644 index 00000000..53c8480d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_paired_packedpartial_gridcap160_0194_v1.py @@ -0,0 +1,135 @@ +"""Flash-KMeans high-D paired packed-partial grid-cap-160 candidate. + +Minimum architecture: sm_100a. This additive bucket-kernel candidate keeps the +Blackwell tcgen05/TMEM packed-key distance producer on the contract-visible +path for paired B=1,N=2048,K=4096 high-D rows. D=512 changes producer work +ownership from the 128-CTA cap to a 160-CTA cap while keeping the validated R2 +reducer; D=448 and other high-D focus rows delegate to the gridcap128 seed +because the 160-CTA cap was slower for D=448. It is not intended for +sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +import os +from typing import Any +from . import flash_kmeans_assign_highd_paired_packedpartial_gridcap128_8512_v1 as _gridcap128 +from . import flash_kmeans_assign_highd_paired_packedpartial_r1_r1c1_v1 as _r1 +from . import flash_kmeans_assign_highd_paired_packedpartial_r2_7b3c_v1 as _r2 +from . import flash_kmeans_assign_highd_splitk_8de8_v1 as _base +from . import flash_kmeans_assign_highd_splitk_blockn64_g2r4_b5a6_v1 as _g2r4 +from . import flash_kmeans_assign_highd_splitk_blockn64_g2r4_streamdep_4f2c_v1 as _streamdep +from .._dispatch_runtime import pack_kernel_args +BLOCK_N = _g2r4.BLOCK_N +MMA_BLOCK_K = _g2r4.MMA_BLOCK_K +SPLITK_GROUP_K_TILES = _g2r4.SPLITK_GROUP_K_TILES +SPLITK_TILE_K = _g2r4.SPLITK_TILE_K +SPLITK_GRID_CAP = _g2r4.SPLITK_GRID_CAP +PRODUCER_GRID_CAP = 160 +SUPPORTED_DIMS = set(_base.SUPPORTED_DIMS) +BF16_DTYPE_NAMES = set(_base.BF16_DTYPE_NAMES) +PAIRED_PACKEDPARTIAL_DIMS = {448, 512} +ROUTE_ID = 'highd_paired_packedpartial_gridcap160_0194_v1' +SEED_ID = 'highd-paired-packedpartial-gridcap160-0194-v1' + +def _verify_export_ir() -> Any: + verify_kernel = os.environ.get('LOOM_FLASH_KMEANS_HIGHD_PAIRED_PACKEDPARTIAL_GRIDCAP160_0194_VERIFY_KERNEL') + if verify_kernel == 'r1_reduce': + return _r1.reduce_ir + if verify_kernel == 'r2_partial': + return _r2.partial_ir + if verify_kernel == 'r2_reduce': + return _r2.reduce_ir + if verify_kernel == 'streamdep_partial': + return _streamdep.partial_ir + if verify_kernel == 'streamdep_reduce': + return _streamdep.reduce_ir + return _r1.partial_ir +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_packedpartial_producer_r1c1_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_keys", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 192}')) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + _validate_shape(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters) + if _use_paired_gridcap160(bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters): + child = _r2 + reducer = 'r2_two_lane' + _launch_paired_gridcap(inputs, child=child, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters) + trace = _route_trace_paired(inputs, child=child, reducer=reducer) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': trace} + outputs = _gridcap128.launch_for_eval(inputs) + trace = _wrap_child_trace(inputs, child_outputs=outputs) + return {'cluster_ids': outputs['cluster_ids'], 'selected_route': ROUTE_ID, 'route_trace': trace} + +def _use_paired_gridcap160(*, bsz: int, n_points: int, dim: int, n_clusters: int) -> bool: + return bsz == 1 and n_points == 2048 and (n_clusters == 4096) and (dim == 512) + +def _validate_shape(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int) -> None: + del bsz + dtype_name = _base._dtype_name(inputs) + if dtype_name not in BF16_DTYPE_NAMES: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires bfloat16 input, got ', format(dtype_name, '')])) + if dim not in SUPPORTED_DIMS: + raise ValueError(''.join([format(ROUTE_ID, ''), ' supports D=', format(sorted(SUPPORTED_DIMS), ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % MMA_BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by MMA_BLOCK_K=', format(MMA_BLOCK_K, ''), ', got ', format(n_clusters, '')])) + +def _launch_paired_gridcap(inputs: dict[str, Any], *, child: Any, bsz: int, n_points: int, dim: int, n_clusters: int) -> None: + import torch + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // MMA_BLOCK_K + k_slices = k_tiles // SPLITK_GROUP_K_TILES + total_work = bsz * num_n_tiles * k_slices + partial_keys = child._partial_key_buffer(inputs, total_work) + tmap_x, tmap_c = _g2r4._make_tmaps(inputs) + stream = torch.cuda.current_stream() + partial_kernel, smem_bytes, threads = child._loaded_partial_key_kernel() + partial_args = pack_kernel_args(child.partial_ir, x_tmap=tmap_x, c_tmap=tmap_c, c_sq=inputs['c_sq'], partial_keys=partial_keys, B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles, K_slices=k_slices) + partial_kernel.launch(grid=(min(total_work, PRODUCER_GRID_CAP), 1, 1), block=(threads, 1, 1), args=partial_args, shared_mem=smem_bytes, stream=stream) + reduce_kernel, smem_bytes, threads = child._loaded_reduce_key_kernel() + reduce_args = [partial_keys, inputs['out'], bsz, n_points, n_clusters, num_n_tiles, k_slices] + reduce_kernel.launch(grid=(min(bsz * num_n_tiles, SPLITK_GRID_CAP), 1, 1), block=(threads, 1, 1), args=reduce_args, shared_mem=smem_bytes, stream=stream) + +def _route_trace_paired(inputs: dict[str, Any], *, child: Any, reducer: str) -> dict[str, Any]: + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // MMA_BLOCK_K + k_slices = k_tiles // SPLITK_GROUP_K_TILES + total_work = int(inputs['B']) * num_n_tiles * k_slices + producer_grid = min(total_work, PRODUCER_GRID_CAP) + child_entrypoint = 'loom.examples.weave.flash_kmeans_assign_highd_paired_packedpartial_r1_r1c1_v1:launch_for_eval' if child is _r1 else 'loom.examples.weave.flash_kmeans_assign_highd_paired_packedpartial_r2_7b3c_v1:launch_for_eval' + child_route = _r1.ROUTE_ID if child is _r1 else _r2.ROUTE_ID + child_guard = 'guard_highd_paired_packedpartial_r1_r1c1_v1' if child is _r1 else 'guard_highd_paired_packedpartial_r2_7b3c_v1' + return {'shape_key': _base._shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_highd_paired_packedpartial_gridcap160_0194_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'shape-specific-seed-composite', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_highd_paired_packedpartial_gridcap160_0194_v1', 'guard_condition': 'dtype == bfloat16 and B == 1 and N == 2048 and K == 4096 and D == 512; D=512 uses packed partial R2 with producer grid capped at 160 CTAs', 'classification': 'route-ok', 'child_route': child_route, 'child_entrypoint': child_entrypoint, 'child_guard_id': child_guard, 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'total_tiles': total_work, 'launch_grid': (producer_grid, 1, 1), 'paired_packedpartial_route': True, 'selected_reducer_variant': reducer, 'tile_shape': {'BLOCK_N': BLOCK_N, 'BLOCK_K': SPLITK_TILE_K, 'mma_block_k': MMA_BLOCK_K, 'split_k_slices': k_slices, 'grouped_mma_k_tiles': SPLITK_GROUP_K_TILES, 'producer_work_items': total_work, 'producer_grid_cap': PRODUCER_GRID_CAP, 'producer_grid_work_items_per_cta_floor': total_work // producer_grid, 'producer_grid_remainder_work_items': total_work % producer_grid, 'producer_to_consumer': 'u64_partial_key_buffer_plus_shape_guarded_r1_or_r2_reduce', 'partial_key': 'ordered_f32_score_high32_inverse_cluster_index_low32', 'reducer_lanes_per_row': _r1.REDUCE_LANES_PER_ROW if child is _r1 else _r2.REDUCE_LANES_PER_ROW, 'producer_to_reducer_sync': 'same_stream_ordering'}, 'reason': 'paired D512 row keeps the G2 tcgen05 packed-key producer but caps producer CTAs at a small over-subscription of the B200 SM count to test one full producer wave plus a residual wave'} + +def _wrap_child_trace(inputs: dict[str, Any], *, child_outputs: dict[str, Any]) -> dict[str, Any]: + child_trace = dict(child_outputs.get('route_trace', {})) + trace = dict(child_trace) + trace.update({'shape_key': _base._shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_highd_paired_packedpartial_gridcap160_0194_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'shape-specific-seed-composite', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_highd_paired_packedpartial_gridcap160_0194_v1', 'guard_condition': 'dtype == bfloat16 and D in [384,448,512] and N % 64 == 0 and K % 256 == 0; paired B=1,N=2048,K=4096,D=512 uses 160-grid-cap packed partial keys, D=448 and other high-D rows delegate to highd_paired_packedpartial_gridcap128_8512_v1', 'classification': 'route-ok', 'child_route': child_outputs.get('selected_route') or child_trace.get('selected_route'), 'child_entrypoint': child_trace.get('selected_entrypoint'), 'child_guard_id': child_trace.get('guard_id'), 'child_tile_shape': child_trace.get('tile_shape'), 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'paired_packedpartial_route': bool(child_trace.get('paired_packedpartial_route', False)), 'reason': 'D448 and non-paired high-D focus rows delegate to the prior gridcap128 composite seed'}) + return trace + +def compile_and_launch_flash_kmeans_assign_highd_paired_packedpartial(B: int=1, N: int=2048, K: int=4096, D: int=512, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + seed = 44801 if D == 448 else 51202 + torch.manual_seed(seed) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': seed}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_paired_packedpartial_r1_r1c1_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_paired_packedpartial_r1_r1c1_v1.py new file mode 100644 index 00000000..25e0bbc4 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_paired_packedpartial_r1_r1c1_v1.py @@ -0,0 +1,159 @@ +"""Flash-KMeans high-D paired packed-partial R1 Split-K candidate. + +Minimum architecture: sm_100a. This additive bucket-kernel candidate keeps the +Blackwell tcgen05/TMEM G2 distance-product producer on the contract-visible +path for paired B=1,N=2048,K=4096 high-D rows, but replaces separate +partial-score and partial-index buffers with one ordered u64 key per row and +K-slice. A single reducer thread scans the packed keys for each row and writes +cluster IDs. +Huge-K focus rows delegate to the validated 4f2c stream-ordered child. It is +not intended for sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +import os +from functools import lru_cache +from typing import Any +from . import flash_kmeans_assign_highd_splitk_8de8_v1 as _base +from . import flash_kmeans_assign_highd_splitk_blockn64_g2r4_b5a6_v1 as _g2r4 +from . import flash_kmeans_assign_highd_splitk_blockn64_g2r4_streamdep_4f2c_v1 as _streamdep +from .flash_kmeans_assign_stream import stream_cache_key +from .._dispatch_runtime import pack_kernel_args +BLOCK_N = _g2r4.BLOCK_N +MMA_BLOCK_K = _g2r4.MMA_BLOCK_K +FEAT_CHUNK = _g2r4.FEAT_CHUNK +ROW16_LOAD_COUNT = _g2r4.ROW16_LOAD_COUNT +ROW16_LOGICAL_CHUNK_K = _g2r4.ROW16_LOGICAL_CHUNK_K +CSQ_STAGE_VEC = _g2r4.CSQ_STAGE_VEC +X_TILE_BYTES = _g2r4.X_TILE_BYTES +C_TILE_BYTES = _g2r4.C_TILE_BYTES +CSQ_TILE_BYTES = _g2r4.CSQ_TILE_BYTES +SPLITK_GROUP_K_TILES = _g2r4.SPLITK_GROUP_K_TILES +SPLITK_TILE_K = _g2r4.SPLITK_TILE_K +SPLITK_GRID_CAP = _g2r4.SPLITK_GRID_CAP +SUPPORTED_DIMS = set(_base.SUPPORTED_DIMS) +BF16_DTYPE_NAMES = set(_base.BF16_DTYPE_NAMES) +PAIRED_PACKEDPARTIAL_DIMS = {448, 512} +NUM_COMPUTE_WARPS = _g2r4.NUM_COMPUTE_WARPS +REDUCE_LANES_PER_ROW = 1 +REDUCE_THREADS = BLOCK_N * REDUCE_LANES_PER_ROW +PAIRED_K_SLICES = 8 +U32_MASK = 4294967295 +ROUTE_ID = 'highd_paired_packedpartial_r1_r1c1_v1' +SEED_ID = 'highd-paired-packedpartial-r1-r1c1-v1' +flash_kmeans_assign_highd_paired_packedpartial_producer_r1c1_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_packedpartial_producer_r1c1_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_keys", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 192}')) +flash_kmeans_assign_highd_paired_packedpartial_reduce_r1_r1c1_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_packedpartial_reduce_r1_r1c1_v1", "arg_keys": ["partial_keys", "out", "B", "N", "K", "num_n_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 64}')) +partial_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_packedpartial_producer_r1c1_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_keys", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 192}')) +reduce_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_packedpartial_reduce_r1_r1c1_v1", "arg_keys": ["partial_keys", "out", "B", "N", "K", "num_n_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 64}')) + +def _verify_export_ir() -> Any: + verify_kernel = os.environ.get('LOOM_FLASH_KMEANS_HIGHD_PAIRED_PACKEDPARTIAL_R1_R1C1_VERIFY_KERNEL') + if verify_kernel == 'reduce': + return reduce_ir + if verify_kernel == 'streamdep_partial': + return _streamdep.partial_ir + if verify_kernel == 'streamdep_reduce': + return _streamdep.reduce_ir + return partial_ir +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_packedpartial_producer_r1c1_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_keys", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 192}')) + +@lru_cache(maxsize=1) +def _loaded_partial_key_kernel() -> tuple[Any, int, int]: + from .._dispatch_runtime import CUDAKernel + cubin, kernel_name, smem_bytes, threads = _base._compile_ir(partial_ir) + return (CUDAKernel(cubin, kernel_name), smem_bytes, threads) + +@lru_cache(maxsize=1) +def _loaded_reduce_key_kernel() -> tuple[Any, int, int]: + from .._dispatch_runtime import CUDAKernel + cubin, kernel_name, smem_bytes, threads = _base._compile_ir(reduce_ir) + return (CUDAKernel(cubin, kernel_name), smem_bytes, threads) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + _validate_shape(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters) + if _use_paired_packedpartial(bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters): + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // MMA_BLOCK_K + k_slices = k_tiles // SPLITK_GROUP_K_TILES + total_work = bsz * num_n_tiles * k_slices + _launch_paired_packedpartial(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters, num_n_tiles=num_n_tiles, k_tiles=k_tiles, k_slices=k_slices, total_work=total_work) + trace = _route_trace_paired(inputs, total_work=total_work, grid=(min(total_work, SPLITK_GRID_CAP), 1, 1), k_slices=k_slices) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': trace} + outputs = _streamdep.launch_for_eval(inputs) + trace = _wrap_child_trace(inputs, child_outputs=outputs) + return {'cluster_ids': outputs['cluster_ids'], 'selected_route': ROUTE_ID, 'route_trace': trace} + +def _use_paired_packedpartial(*, bsz: int, n_points: int, dim: int, n_clusters: int) -> bool: + return bsz == 1 and n_points == 2048 and (n_clusters == 4096) and (dim in PAIRED_PACKEDPARTIAL_DIMS) + +def _validate_shape(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int) -> None: + del bsz + dtype_name = _base._dtype_name(inputs) + if dtype_name not in BF16_DTYPE_NAMES: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires bfloat16 input, got ', format(dtype_name, '')])) + if dim not in SUPPORTED_DIMS: + raise ValueError(''.join([format(ROUTE_ID, ''), ' supports D=', format(sorted(SUPPORTED_DIMS), ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % MMA_BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by MMA_BLOCK_K=', format(MMA_BLOCK_K, ''), ', got ', format(n_clusters, '')])) + +def _launch_paired_packedpartial(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int, num_n_tiles: int, k_tiles: int, k_slices: int, total_work: int) -> None: + import torch + if k_slices != PAIRED_K_SLICES: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires K_slices=', format(PAIRED_K_SLICES, ''), ', got ', format(k_slices, '')])) + partial_keys = _partial_key_buffer(inputs, total_work) + tmap_x, tmap_c = _g2r4._make_tmaps(inputs) + stream = torch.cuda.current_stream() + partial_kernel, smem_bytes, threads = _loaded_partial_key_kernel() + partial_args = pack_kernel_args(partial_ir, x_tmap=tmap_x, c_tmap=tmap_c, c_sq=inputs['c_sq'], partial_keys=partial_keys, B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles, K_slices=k_slices) + partial_kernel.launch(grid=(min(total_work, SPLITK_GRID_CAP), 1, 1), block=(threads, 1, 1), args=partial_args, shared_mem=smem_bytes, stream=stream) + reduce_kernel, smem_bytes, threads = _loaded_reduce_key_kernel() + reduce_args = [partial_keys, inputs['out'], bsz, n_points, n_clusters, num_n_tiles, k_slices] + reduce_kernel.launch(grid=(min(bsz * num_n_tiles, SPLITK_GRID_CAP), 1, 1), block=(threads, 1, 1), args=reduce_args, shared_mem=smem_bytes, stream=stream) + +def _partial_key_buffer(inputs: dict[str, Any], total_work: int) -> Any: + import torch + cache_key = stream_cache_key(inputs, int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), int(inputs['D']), int(total_work), int(BLOCK_N), int(SPLITK_GROUP_K_TILES)) + cache = inputs.setdefault('_flash_kmeans_assign_highd_paired_packedpartial_r1_r1c1_v1_keys', {}) + cached = cache.get(cache_key) + if cached is not None: + return cached + keys = torch.empty((total_work, BLOCK_N), dtype=torch.int64, device=inputs['x'].device) + cache[cache_key] = keys + return keys + +def _route_trace_paired(inputs: dict[str, Any], *, total_work: int, grid: tuple[int, int, int], k_slices: int) -> dict[str, Any]: + return {'shape_key': _base._shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_highd_paired_packedpartial_r1_r1c1_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'shape-specific-seed-composite', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_highd_paired_packedpartial_r1_r1c1_v1', 'guard_condition': 'dtype == bfloat16 and B == 1 and N == 2048 and K == 4096 and D in [448, 512]', 'classification': 'route-ok', 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'total_tiles': total_work, 'launch_grid': grid, 'tile_shape': {'BLOCK_N': BLOCK_N, 'BLOCK_K': SPLITK_TILE_K, 'mma_block_k': MMA_BLOCK_K, 'split_k_slices': k_slices, 'grouped_mma_k_tiles': SPLITK_GROUP_K_TILES, 'producer_work_items': total_work, 'producer_to_consumer': 'u64_partial_key_buffer_plus_r1_reduce', 'partial_key': 'ordered_f32_score_high32_inverse_cluster_index_low32', 'reducer_lanes_per_row': REDUCE_LANES_PER_ROW, 'reducer_fixed_k_slices': PAIRED_K_SLICES, 'producer_to_reducer_sync': 'same_stream_ordering'}, 'reason': 'paired high-D row keeps the G2 tcgen05 Split-K producer but packs score/index partials into one u64 buffer and scans all K-slice keys with a single row thread'} + +def _wrap_child_trace(inputs: dict[str, Any], *, child_outputs: dict[str, Any]) -> dict[str, Any]: + child_trace = dict(child_outputs.get('route_trace', {})) + child_route = child_outputs.get('selected_route') or child_trace.get('selected_route') + trace = dict(child_trace) + trace.update({'shape_key': _base._shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_highd_paired_packedpartial_r1_r1c1_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'shape-specific-seed-composite', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_highd_paired_packedpartial_r1_r1c1_v1', 'guard_condition': 'dtype == bfloat16 and D in [384,448,512] and N % 64 == 0 and K % 256 == 0; paired B=1,N=2048,K=4096,D in [448,512] uses packed partial keys, all other focus rows use highd_splitk_blockn64_g2r4_streamdep_4f2c_v1', 'classification': 'route-ok', 'child_route': child_route, 'child_entrypoint': child_trace.get('selected_entrypoint'), 'child_guard_id': child_trace.get('guard_id'), 'child_tile_shape': child_trace.get('tile_shape'), 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'reason': 'non-paired high-D focus row keeps the stream-ordered G2/R4 Split-K child'}) + return trace + +def compile_and_launch_flash_kmeans_assign_highd_paired_packedpartial(B: int=1, N: int=2048, K: int=4096, D: int=512, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(51202) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 51202}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_paired_packedpartial_r2_7b3c_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_paired_packedpartial_r2_7b3c_v1.py new file mode 100644 index 00000000..32c1d4d8 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_paired_packedpartial_r2_7b3c_v1.py @@ -0,0 +1,170 @@ +"""Flash-KMeans high-D paired packed-partial R2 Split-K candidate. + +Minimum architecture: sm_100a. This additive bucket-kernel candidate keeps the +Blackwell tcgen05/TMEM G2 distance-product producer on the contract-visible +path for paired B=1,N=2048,K=4096 high-D rows, but replaces separate +partial-score and partial-index buffers with one ordered u64 key per row and +K-slice. A two-lane reducer scans packed keys and writes cluster IDs. +Huge-K focus rows delegate to the validated 4f2c stream-ordered child. It is +not intended for sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +import os +from functools import lru_cache +from typing import Any +from . import flash_kmeans_assign_highd_splitk_8de8_v1 as _base +from . import flash_kmeans_assign_highd_splitk_blockn64_g2r4_b5a6_v1 as _g2r4 +from . import flash_kmeans_assign_highd_splitk_blockn64_g2r4_streamdep_4f2c_v1 as _streamdep +from .flash_kmeans_assign_stream import stream_cache_key +from .._dispatch_runtime import pack_kernel_args +BLOCK_N = _g2r4.BLOCK_N +MMA_BLOCK_K = _g2r4.MMA_BLOCK_K +FEAT_CHUNK = _g2r4.FEAT_CHUNK +ROW16_LOAD_COUNT = _g2r4.ROW16_LOAD_COUNT +ROW16_LOGICAL_CHUNK_K = _g2r4.ROW16_LOGICAL_CHUNK_K +CSQ_STAGE_VEC = _g2r4.CSQ_STAGE_VEC +X_TILE_BYTES = _g2r4.X_TILE_BYTES +C_TILE_BYTES = _g2r4.C_TILE_BYTES +CSQ_TILE_BYTES = _g2r4.CSQ_TILE_BYTES +DEFAULT_SPLITK_GROUP_K_TILES = _g2r4.SPLITK_GROUP_K_TILES +SPLITK_GROUP_K_TILES_ENV = 'LOOM_FLASH_KMEANS_HIGHD_PAIRED_PACKEDPARTIAL_R2_7B3C_GROUP_K_TILES' + +def _configured_splitk_group_k_tiles() -> int: + raw = os.environ.get(SPLITK_GROUP_K_TILES_ENV) + if raw is None or raw == '': + return DEFAULT_SPLITK_GROUP_K_TILES + value = int(raw) + if value not in (1, DEFAULT_SPLITK_GROUP_K_TILES): + raise ValueError(''.join([format(SPLITK_GROUP_K_TILES_ENV, ''), ' must be 1 or ', format(DEFAULT_SPLITK_GROUP_K_TILES, ''), ', got ', format(value, '')])) + return value +SPLITK_GROUP_K_TILES = _decode_capture(_json_loads('2')) +SPLITK_TILE_K = MMA_BLOCK_K * SPLITK_GROUP_K_TILES +SPLITK_GRID_CAP = _g2r4.SPLITK_GRID_CAP +SUPPORTED_DIMS = set(_base.SUPPORTED_DIMS) +BF16_DTYPE_NAMES = set(_base.BF16_DTYPE_NAMES) +PAIRED_PACKEDPARTIAL_DIMS = {448, 512} +NUM_COMPUTE_WARPS = _g2r4.NUM_COMPUTE_WARPS +REDUCE_LANES_PER_ROW = 2 +REDUCE_THREADS = BLOCK_N * REDUCE_LANES_PER_ROW +U32_MASK = 4294967295 +if SPLITK_GROUP_K_TILES == DEFAULT_SPLITK_GROUP_K_TILES: + ROUTE_ID = 'highd_paired_packedpartial_r2_7b3c_v1' + SEED_ID = 'highd-paired-packedpartial-r2-7b3c-v1' +else: + ROUTE_ID = 'highd_paired_packedpartial_r2_g1_tileprobe_7b3c_v1' + SEED_ID = 'highd-paired-packedpartial-r2-g1-tileprobe-7b3c-v1' +flash_kmeans_assign_highd_paired_packedpartial_producer_7b3c_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_packedpartial_producer_7b3c_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_keys", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 192}')) +flash_kmeans_assign_highd_paired_packedpartial_reduce_r2_7b3c_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_packedpartial_reduce_r2_7b3c_v1", "arg_keys": ["partial_keys", "out", "B", "N", "K", "num_n_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 128}')) +partial_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_packedpartial_producer_7b3c_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_keys", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 192}')) +reduce_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_packedpartial_reduce_r2_7b3c_v1", "arg_keys": ["partial_keys", "out", "B", "N", "K", "num_n_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 128}')) + +def _verify_export_ir() -> Any: + verify_kernel = os.environ.get('LOOM_FLASH_KMEANS_HIGHD_PAIRED_PACKEDPARTIAL_R2_7B3C_VERIFY_KERNEL') + if verify_kernel == 'reduce': + return reduce_ir + if verify_kernel == 'streamdep_partial': + return _streamdep.partial_ir + if verify_kernel == 'streamdep_reduce': + return _streamdep.reduce_ir + return partial_ir +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_packedpartial_producer_7b3c_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_keys", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 192}')) + +@lru_cache(maxsize=1) +def _loaded_partial_key_kernel() -> tuple[Any, int, int]: + from .._dispatch_runtime import CUDAKernel + cubin, kernel_name, smem_bytes, threads = _base._compile_ir(partial_ir) + return (CUDAKernel(cubin, kernel_name), smem_bytes, threads) + +@lru_cache(maxsize=1) +def _loaded_reduce_key_kernel() -> tuple[Any, int, int]: + from .._dispatch_runtime import CUDAKernel + cubin, kernel_name, smem_bytes, threads = _base._compile_ir(reduce_ir) + return (CUDAKernel(cubin, kernel_name), smem_bytes, threads) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + _validate_shape(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters) + if _use_paired_packedpartial(bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters): + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // MMA_BLOCK_K + k_slices = k_tiles // SPLITK_GROUP_K_TILES + total_work = bsz * num_n_tiles * k_slices + _launch_paired_packedpartial(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters, num_n_tiles=num_n_tiles, k_tiles=k_tiles, k_slices=k_slices, total_work=total_work) + trace = _route_trace_paired(inputs, total_work=total_work, grid=(min(total_work, SPLITK_GRID_CAP), 1, 1), k_slices=k_slices) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': trace} + outputs = _streamdep.launch_for_eval(inputs) + trace = _wrap_child_trace(inputs, child_outputs=outputs) + return {'cluster_ids': outputs['cluster_ids'], 'selected_route': ROUTE_ID, 'route_trace': trace} + +def _use_paired_packedpartial(*, bsz: int, n_points: int, dim: int, n_clusters: int) -> bool: + return bsz == 1 and n_points == 2048 and (n_clusters == 4096) and (dim in PAIRED_PACKEDPARTIAL_DIMS) + +def _validate_shape(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int) -> None: + del bsz + dtype_name = _base._dtype_name(inputs) + if dtype_name not in BF16_DTYPE_NAMES: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires bfloat16 input, got ', format(dtype_name, '')])) + if dim not in SUPPORTED_DIMS: + raise ValueError(''.join([format(ROUTE_ID, ''), ' supports D=', format(sorted(SUPPORTED_DIMS), ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % MMA_BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by MMA_BLOCK_K=', format(MMA_BLOCK_K, ''), ', got ', format(n_clusters, '')])) + +def _launch_paired_packedpartial(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int, num_n_tiles: int, k_tiles: int, k_slices: int, total_work: int) -> None: + import torch + partial_keys = _partial_key_buffer(inputs, total_work) + tmap_x, tmap_c = _g2r4._make_tmaps(inputs) + stream = torch.cuda.current_stream() + partial_kernel, smem_bytes, threads = _loaded_partial_key_kernel() + partial_args = pack_kernel_args(partial_ir, x_tmap=tmap_x, c_tmap=tmap_c, c_sq=inputs['c_sq'], partial_keys=partial_keys, B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles, K_slices=k_slices) + partial_kernel.launch(grid=(min(total_work, SPLITK_GRID_CAP), 1, 1), block=(threads, 1, 1), args=partial_args, shared_mem=smem_bytes, stream=stream) + reduce_kernel, smem_bytes, threads = _loaded_reduce_key_kernel() + reduce_args = [partial_keys, inputs['out'], bsz, n_points, n_clusters, num_n_tiles, k_slices] + reduce_kernel.launch(grid=(min(bsz * num_n_tiles, SPLITK_GRID_CAP), 1, 1), block=(threads, 1, 1), args=reduce_args, shared_mem=smem_bytes, stream=stream) + +def _partial_key_buffer(inputs: dict[str, Any], total_work: int) -> Any: + import torch + cache_key = stream_cache_key(inputs, int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), int(inputs['D']), int(total_work), int(BLOCK_N), int(SPLITK_GROUP_K_TILES)) + cache = inputs.setdefault('_flash_kmeans_assign_highd_paired_packedpartial_r2_7b3c_v1_keys', {}) + cached = cache.get(cache_key) + if cached is not None: + return cached + keys = torch.empty((total_work, BLOCK_N), dtype=torch.int64, device=inputs['x'].device) + cache[cache_key] = keys + return keys + +def _route_trace_paired(inputs: dict[str, Any], *, total_work: int, grid: tuple[int, int, int], k_slices: int) -> dict[str, Any]: + return {'shape_key': _base._shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_highd_paired_packedpartial_r2_7b3c_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'shape-specific-seed-composite', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_highd_paired_packedpartial_r2_7b3c_v1', 'guard_condition': ''.join(['dtype == bfloat16 and B == 1 and N == 2048 and K == 4096 and D in [448, 512]; grouped_mma_k_tiles=', format(SPLITK_GROUP_K_TILES, '')]), 'classification': 'route-ok', 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'total_tiles': total_work, 'launch_grid': grid, 'tile_shape': {'BLOCK_N': BLOCK_N, 'BLOCK_K': SPLITK_TILE_K, 'mma_block_k': MMA_BLOCK_K, 'split_k_slices': k_slices, 'grouped_mma_k_tiles': SPLITK_GROUP_K_TILES, 'producer_work_items': total_work, 'producer_to_consumer': 'u64_partial_key_buffer_plus_r2_reduce', 'partial_key': 'ordered_f32_score_high32_inverse_cluster_index_low32', 'reducer_lanes_per_row': REDUCE_LANES_PER_ROW, 'producer_to_reducer_sync': 'same_stream_ordering', 'tile_probe_env': SPLITK_GROUP_K_TILES_ENV if SPLITK_GROUP_K_TILES != DEFAULT_SPLITK_GROUP_K_TILES else None}, 'reason': ''.join(['paired high-D row keeps the G2 tcgen05 Split-K producer but packs score/index partials into one u64 buffer and scans keys with a two-lane row reducer; grouped_mma_k_tiles=', format(SPLITK_GROUP_K_TILES, '')])} + +def _wrap_child_trace(inputs: dict[str, Any], *, child_outputs: dict[str, Any]) -> dict[str, Any]: + child_trace = dict(child_outputs.get('route_trace', {})) + child_route = child_outputs.get('selected_route') or child_trace.get('selected_route') + trace = dict(child_trace) + trace.update({'shape_key': _base._shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_highd_paired_packedpartial_r2_7b3c_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'shape-specific-seed-composite', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_highd_paired_packedpartial_r2_7b3c_v1', 'guard_condition': 'dtype == bfloat16 and D in [384,448,512] and N % 64 == 0 and K % 256 == 0; paired B=1,N=2048,K=4096,D in [448,512] uses packed partial keys, all other focus rows use highd_splitk_blockn64_g2r4_streamdep_4f2c_v1', 'classification': 'route-ok', 'child_route': child_route, 'child_entrypoint': child_trace.get('selected_entrypoint'), 'child_guard_id': child_trace.get('guard_id'), 'child_tile_shape': child_trace.get('tile_shape'), 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'reason': 'non-paired high-D focus row keeps the stream-ordered G2/R4 Split-K child'}) + return trace + +def compile_and_launch_flash_kmeans_assign_highd_paired_packedpartial(B: int=1, N: int=2048, K: int=4096, D: int=512, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(51202) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 51202}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_paired_xreuse_dualtmem_r47_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_paired_xreuse_dualtmem_r47_v1.py new file mode 100644 index 00000000..41c08c7f --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_paired_xreuse_dualtmem_r47_v1.py @@ -0,0 +1,137 @@ +"""Flash-KMeans high-D paired D448 X-reuse dual-TMEM round-47 candidate. + +Minimum architecture: sm_100a. This additive exact-bucket seed keeps the +Blackwell tcgen05/TMEM packed-key producer on the contract-visible path for +B=1,N=2048,K=4096,D=448, but reorders the paired local-K producer so each +64-wide X feature tile is loaded once and consumed by two 256-column TMEM +accumulators. The D448 reducer remains the round-39 unrolled R1 reducer, and +D512 delegates to the round-39 route. It is not intended for sm_120a/sm_121a +where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +import os +from functools import lru_cache +from typing import Any +from . import flash_kmeans_assign_highd_paired_ownerreduce_r39_v1 as _r39 +from . import flash_kmeans_assign_highd_paired_packedpartial_r1_r1c1_v1 as _r1 +from . import flash_kmeans_assign_highd_splitk_8de8_v1 as _base +from . import flash_kmeans_assign_highd_splitk_blockn64_g2r4_b5a6_v1 as _g2r4 +from .._dispatch_runtime import pack_kernel_args +BLOCK_N = _r39.BLOCK_N +MMA_BLOCK_K = _r39.MMA_BLOCK_K +FEAT_CHUNK = _g2r4.FEAT_CHUNK +ROW16_LOAD_COUNT = _g2r4.ROW16_LOAD_COUNT +ROW16_LOGICAL_CHUNK_K = _g2r4.ROW16_LOGICAL_CHUNK_K +CSQ_STAGE_VEC = _g2r4.CSQ_STAGE_VEC +X_TILE_BYTES = _g2r4.X_TILE_BYTES +C_TILE_BYTES = _g2r4.C_TILE_BYTES +CSQ_TILE_BYTES = _g2r4.CSQ_TILE_BYTES +SPLITK_GROUP_K_TILES = _r39.SPLITK_GROUP_K_TILES +C_PAIR_BYTES = SPLITK_GROUP_K_TILES * C_TILE_BYTES +SPLITK_TILE_K = _r39.SPLITK_TILE_K +SPLITK_GRID_CAP = _r39.SPLITK_GRID_CAP +BF16_DTYPE_NAMES = set(_base.BF16_DTYPE_NAMES) +PRODUCER_GRID_CAP_D448 = _r39.PRODUCER_GRID_CAP_D448 +PAIRED_K_SLICES = _r39.PAIRED_K_SLICES +D448_FEATURE_TILES = 448 // FEAT_CHUNK +NUM_COMPUTE_WARPS = _g2r4.NUM_COMPUTE_WARPS +U32_MASK = _r39.U32_MASK +ROUTE_ID = 'highd_paired_xreuse_dualtmem_r47_v1' +SEED_ID = 'highd-paired-xreuse-dualtmem-r47-v1' +flash_kmeans_assign_highd_paired_xreuse_dualtmem_producer_r47_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_xreuse_dualtmem_producer_r47_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_keys", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 75776, "constants": [], "cta_group": 1, "threads": 192}')) +partial_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_xreuse_dualtmem_producer_r47_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_keys", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 75776, "constants": [], "cta_group": 1, "threads": 192}')) +reduce_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_ownerreduce_r39_reduce1_unroll_v1", "arg_keys": ["partial_keys", "out", "B", "N", "K", "num_n_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 64}')) + +def _verify_export_ir() -> Any: + verify_kernel = os.environ.get('LOOM_FLASH_KMEANS_HIGHD_PAIRED_XREUSE_DUALTMEM_R47_VERIFY_KERNEL') + if verify_kernel == 'd448_reduce': + return reduce_ir + if verify_kernel == 'd512_partial': + return _r39._r2.partial_ir + if verify_kernel == 'd512_reduce': + return _r39._r2.reduce_ir + return partial_ir +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_xreuse_dualtmem_producer_r47_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_keys", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 75776, "constants": [], "cta_group": 1, "threads": 192}')) + +@lru_cache(maxsize=1) +def _loaded_partial_key_kernel() -> tuple[Any, int, int]: + from .._dispatch_runtime import CUDAKernel + cubin, kernel_name, smem_bytes, threads = _base._compile_ir(partial_ir) + return (CUDAKernel(cubin, kernel_name), smem_bytes, threads) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + _validate_shape(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters) + if _use_d448_custom(bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters): + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // MMA_BLOCK_K + k_slices = k_tiles // SPLITK_GROUP_K_TILES + if k_slices != PAIRED_K_SLICES: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires K_slices=', format(PAIRED_K_SLICES, ''), ', got ', format(k_slices, '')])) + total_work = bsz * num_n_tiles * k_slices + producer_grid, reducer_grid = _launch_d448(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters, num_n_tiles=num_n_tiles, k_tiles=k_tiles, k_slices=k_slices, total_work=total_work) + trace = _route_trace_d448(inputs, total_work=total_work, producer_grid=producer_grid, reducer_grid=reducer_grid, k_slices=k_slices) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': trace} + outputs = _r39.launch_for_eval(inputs) + trace = _wrap_r39_trace(inputs, child_outputs=outputs) + return {'cluster_ids': outputs['cluster_ids'], 'selected_route': ROUTE_ID, 'route_trace': trace} + +def _use_d448_custom(*, bsz: int, n_points: int, dim: int, n_clusters: int) -> bool: + return bsz == 1 and n_points == 2048 and (n_clusters == 4096) and (dim == 448) + +def _validate_shape(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int) -> None: + dtype_name = _base._dtype_name(inputs) + if dtype_name not in BF16_DTYPE_NAMES: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires bfloat16 input, got ', format(dtype_name, '')])) + if (bsz, n_points, n_clusters, dim) not in {(1, 2048, 4096, 448), (1, 2048, 4096, 512)}: + raise ValueError(''.join([format(ROUTE_ID, ''), ' is exact-shape only for paired D448/D512 rows'])) + +def _launch_d448(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int, num_n_tiles: int, k_tiles: int, k_slices: int, total_work: int) -> tuple[tuple[int, int, int], tuple[int, int, int]]: + import torch + partial_keys = _r1._partial_key_buffer(inputs, total_work) + tmap_x, tmap_c = _g2r4._make_tmaps(inputs) + stream = torch.cuda.current_stream() + partial_kernel, smem_bytes, threads = _loaded_partial_key_kernel() + partial_args = pack_kernel_args(partial_ir, x_tmap=tmap_x, c_tmap=tmap_c, c_sq=inputs['c_sq'], partial_keys=partial_keys, B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles, K_slices=k_slices) + producer_grid = (min(total_work, PRODUCER_GRID_CAP_D448), 1, 1) + partial_kernel.launch(grid=producer_grid, block=(threads, 1, 1), args=partial_args, shared_mem=smem_bytes, stream=stream) + reduce_kernel, smem_bytes, threads = _r39._loaded_reduce1_kernel() + reduce_args = [partial_keys, inputs['out'], bsz, n_points, n_clusters, num_n_tiles, k_slices] + reducer_grid = (min(bsz * num_n_tiles, SPLITK_GRID_CAP), 1, 1) + reduce_kernel.launch(grid=reducer_grid, block=(threads, 1, 1), args=reduce_args, shared_mem=smem_bytes, stream=stream) + return (producer_grid, reducer_grid) + +def _route_trace_d448(inputs: dict[str, Any], *, total_work: int, producer_grid: tuple[int, int, int], reducer_grid: tuple[int, int, int], k_slices: int) -> dict[str, Any]: + return {'shape_key': _base._shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_highd_paired_xreuse_dualtmem_r47_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'shape-specific-seed', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_highd_paired_xreuse_dualtmem_r47_v1', 'guard_condition': 'dtype == bfloat16 and B == 1 and N == 2048 and K == 4096 and D in [448,512]; D448 uses X-reuse dual-TMEM producer, D512 delegates to R39', 'classification': 'route-ok', 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'total_tiles': total_work, 'launch_grid': producer_grid, 'reduce_grid': reducer_grid, 'paired_packedpartial_route': True, 'selected_reducer_variant': 'r39_reduce1_unroll8', 'tile_shape': {'BLOCK_N': BLOCK_N, 'BLOCK_K': SPLITK_TILE_K, 'mma_block_k': MMA_BLOCK_K, 'split_k_slices': k_slices, 'feature_tiles': D448_FEATURE_TILES, 'feature_loop_bound': 'compile_time_7', 'feature_loop_unroll': D448_FEATURE_TILES, 'grouped_mma_k_tiles': SPLITK_GROUP_K_TILES, 'x_tma_loads_per_work_item': D448_FEATURE_TILES, 'score_tmem_regions': 2, 'score_tmem_cols': [0, MMA_BLOCK_K], 'producer_work_items': total_work, 'producer_grid_cap': PRODUCER_GRID_CAP_D448, 'producer_grid_work_items_per_cta': total_work // min(total_work, PRODUCER_GRID_CAP_D448), 'producer_to_consumer': 'dual_tmem_xreuse_u64_partial_key_buffer_plus_reduce1_unroll8', 'partial_key': 'ordered_f32_score_high32_inverse_cluster_index_low32', 'reducer_lanes_per_row': 1, 'reducer_threads': _r39.REDUCE_THREADS_R1_D448, 'producer_to_reducer_sync': 'same_stream_ordering', 'reducer_scan_unroll': PAIRED_K_SLICES}, 'reason': 'D448 keeps the R39 packed-key producer/reducer ABI but loads each X feature tile once and accumulates the two paired local-K halves into separate TMEM regions before draining'} + +def _wrap_r39_trace(inputs: dict[str, Any], *, child_outputs: dict[str, Any]) -> dict[str, Any]: + child_trace = dict(child_outputs.get('route_trace', {})) + trace = dict(child_trace) + trace.update({'shape_key': _base._shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_highd_paired_xreuse_dualtmem_r47_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'shape-specific-seed-composite', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_highd_paired_xreuse_dualtmem_r47_v1', 'guard_condition': 'dtype == bfloat16 and B == 1 and N == 2048 and K == 4096 and D in [448,512]; D448 uses X-reuse dual-TMEM producer, D512 delegates to R39', 'classification': 'route-ok', 'child_route': child_outputs.get('selected_route') or child_trace.get('selected_route'), 'child_entrypoint': child_trace.get('selected_entrypoint'), 'child_guard_id': child_trace.get('guard_id'), 'child_tile_shape': child_trace.get('tile_shape'), 'selected_reducer_variant': 'r39_d512_child', 'paired_packedpartial_route': True, 'reason': 'D512 delegates unchanged to the round-39 gridcap160 child route'}) + return trace + +def compile_and_launch_flash_kmeans_assign_highd_paired_xreuse_dualtmem_r47(B: int=1, N: int=2048, K: int=4096, D: int=448, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + seed = 44801 if D == 448 else 51202 + torch.manual_seed(seed) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': seed}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_splitd_6fcf_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_splitd_6fcf_v1.py new file mode 100644 index 00000000..184fb2f6 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_splitd_6fcf_v1.py @@ -0,0 +1,118 @@ +"""Flash-KMeans Euclidean assignment high-D split-D candidate. + +Minimum architecture: sm_100a. This candidate uses Blackwell tcgen05/TMEM +(``lm.mma``) for D in {320, 384, 448, 512}. It slices the feature dimension +into 64-wide TMA/MMA groups so high-D shapes stay within shared memory. It is +not intended for sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +BLOCK_N = 128 +BLOCK_K = 256 +FEAT_CHUNK = 64 +SCORE_CHUNK_K = 128 +CSQ_VEC = 4 +CSQ_STAGE_VEC = 4 +NUM_COMPUTE_WARPS = 4 +X_TILE_BYTES = BLOCK_N * FEAT_CHUNK * 2 +C_TILE_BYTES = BLOCK_K * FEAT_CHUNK * 2 +CSQ_TILE_BYTES = BLOCK_K * 4 +PAIRED_GRID_CAP = 2516 +SUPPORTED_DIMS = {320, 384, 448, 512} +BF16_DTYPE_NAMES = {'bfloat16', 'bf16', 'torch.bfloat16'} +ROUTE_ID = 'highd_splitd_single_tile_6fcf_v1' +SEED_ID = 'highd-splitd-single-tile-6fcf-v1' +flash_kmeans_assign_highd_splitd_6fcf_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitd_6fcf_v1", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 51200, "constants": [], "cta_group": 1, "threads": 192}')) + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as _common_cuda_include_dirs + return _common_cuda_include_dirs() + +def _make_tmaps(inputs: dict[str, Any]) -> tuple[Any, Any]: + cache_key = (int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), int(inputs['D'])) + cached = inputs.get('_flash_kmeans_assign_highd_splitd_6fcf_v1_tmaps') + if cached is not None and cached[0] == cache_key: + return (cached[1], cached[2]) + from .._dispatch_runtime import create_tensor_map_3d + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + dim = int(inputs['D']) + tmap_x = create_tensor_map_3d(inputs['x'].data_ptr(), bsz * n_points, BLOCK_N, dim, FEAT_CHUNK) + tmap_c = create_tensor_map_3d(inputs['centroids'].data_ptr(), bsz * n_clusters, BLOCK_K, dim, FEAT_CHUNK) + inputs['_flash_kmeans_assign_highd_splitd_6fcf_v1_tmaps'] = (cache_key, tmap_x, tmap_c) + return (tmap_x, tmap_c) + +def _compiled_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0327"}, "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", 51200, 192]}')) +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitd_6fcf_v1", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 51200, "constants": [], "cta_group": 1, "threads": 192}')) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + from .._dispatch_runtime import CUDAKernel + from .._dispatch_runtime import pack_kernel_args + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + dtype_name = _dtype_name(inputs) + if dtype_name not in BF16_DTYPE_NAMES: + raise ValueError(''.join(['flash_kmeans_assign_highd_splitd_6fcf_v1 requires bfloat16 input, got ', format(dtype_name, '')])) + if dim not in SUPPORTED_DIMS: + raise ValueError(''.join(['flash_kmeans_assign_highd_splitd_6fcf_v1 supports D=', format(sorted(SUPPORTED_DIMS), ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(n_clusters, '')])) + tmap_x, tmap_c = _make_tmaps(inputs) + cubin, kernel_name, smem_bytes, threads = _compiled_kernel() + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // BLOCK_K + total_tiles = bsz * num_n_tiles + grid = (min(total_tiles, PAIRED_GRID_CAP), 1, 1) + block = (threads, 1, 1) + args = pack_kernel_args(ir, x_tmap=tmap_x, c_tmap=tmap_c, x_sq=inputs['x_sq'], c_sq=inputs['c_sq'], out=inputs['out'], B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles) + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=grid, block=block, args=args, shared_mem=smem_bytes) + trace = _route_trace(inputs, total_tiles=total_tiles, grid=grid) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': trace} + +def compile_and_launch_flash_kmeans_assign_highd_splitd(B: int=1, N: int=1024, K: int=512, D: int=512, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(51201) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 51201}}]) + return result + +def _route_trace(inputs: dict[str, Any], *, total_tiles: int, grid: tuple[int, int, int]) -> dict[str, Any]: + return {'shape_key': _shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_highd_splitd_6fcf_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'specialized', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_highd_splitd_6fcf_v1', 'guard_condition': 'dtype == bfloat16 and D in [320, 384, 448, 512] and N % 128 == 0 and K % 256 == 0', 'classification': 'route-ok', 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'total_tiles': total_tiles, 'launch_grid': grid, 'reason': 'high-D split-D seed accumulates 64-feature tcgen05 slices into one argmin path'} + +def _shape_key(inputs: dict[str, Any]) -> str: + label = inputs.get('label') + if label: + return str(label) + return ''.join(['b', format(int(inputs['B']), ''), '_n', format(int(inputs['N']), ''), '_k', format(int(inputs['K']), ''), '_d', format(int(inputs['D']), '')]) + +def _dtype_name(inputs: dict[str, Any]) -> str: + dtype = inputs.get('dtype') + if dtype is not None: + return str(dtype).replace('torch.', '') + x = inputs.get('x') + if x is not None and hasattr(x, 'dtype'): + return str(x.dtype).replace('torch.', '') + return 'bfloat16' diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_splitk_8de8_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_splitk_8de8_v1.py new file mode 100644 index 00000000..fd78c636 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_splitk_8de8_v1.py @@ -0,0 +1,153 @@ +"""Flash-KMeans Euclidean assignment high-D split-K candidate. + +Minimum architecture: sm_100a. This candidate uses Blackwell tcgen05/TMEM +(``lm.mma``) for D in {320, 384, 448, 512}. High-K rows split the cluster +tile loop across CTAs, write one partial argmin per K tile, then reduce those +partials with a second Weave kernel. Non-splitK rows reuse the validated 6fcf +split-D seed. It is not intended for sm_120a/sm_121a where ptxas rejects +tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +import os +from functools import lru_cache +from typing import Any +from . import flash_kmeans_assign_highd_splitd_6fcf_v1 as _splitd +from .flash_kmeans_assign_stream import stream_cache_key +BLOCK_N = _splitd.BLOCK_N +BLOCK_K = _splitd.BLOCK_K +FEAT_CHUNK = _splitd.FEAT_CHUNK +SCORE_CHUNK_K = _splitd.SCORE_CHUNK_K +CSQ_VEC = _splitd.CSQ_VEC +CSQ_STAGE_VEC = _splitd.CSQ_STAGE_VEC +NUM_COMPUTE_WARPS = _splitd.NUM_COMPUTE_WARPS +X_TILE_BYTES = _splitd.X_TILE_BYTES +C_TILE_BYTES = _splitd.C_TILE_BYTES +CSQ_TILE_BYTES = _splitd.CSQ_TILE_BYTES +SUPPORTED_DIMS = set(_splitd.SUPPORTED_DIMS) +BF16_DTYPE_NAMES = set(_splitd.BF16_DTYPE_NAMES) +SPLITK_GRID_CAP = 4096 +REDUCE_THREADS = 128 +SPLITK_MIN_K_TILES = 16 +ROUTE_ID = 'highd_splitk_8de8_v1' +SEED_ID = 'highd-splitk-8de8-v1' +flash_kmeans_assign_highd_splitk_partial_8de8_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_partial_8de8_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 51200, "constants": [], "cta_group": 1, "threads": 192}')) +flash_kmeans_assign_highd_splitk_reduce_8de8_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_reduce_8de8_v1", "arg_keys": ["partial_scores", "partial_indices", "out", "B", "N", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 128}')) +partial_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_partial_8de8_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 51200, "constants": [], "cta_group": 1, "threads": 192}')) +reduce_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_reduce_8de8_v1", "arg_keys": ["partial_scores", "partial_indices", "out", "B", "N", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 128}')) + +def _verify_export_ir() -> Any: + if os.environ.get('LOOM_FLASH_KMEANS_HIGHD_SPLITK_VERIFY_KERNEL') == 'reduce': + return reduce_ir + return partial_ir +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_partial_8de8_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 51200, "constants": [], "cta_group": 1, "threads": 192}')) + +def _cuda_include_dirs() -> list[str]: + return _splitd._cuda_include_dirs() + +def _compiled_partial_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0325"}, "kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1", 51200, 192]}')) + +def _compiled_reduce_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0326"}, "kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1", 0, 128]}')) + +def _compile_ir(ir_obj: Any) -> tuple[bytes, str, int, int]: + from .._dispatch_runtime import generate_kernel + from .._dispatch_runtime import compile_cuda + smem_bytes = ir_obj.computed_smem_bytes + source = generate_kernel(ir_obj, validate=False, smem_bytes=smem_bytes) + cubin = compile_cuda(source, options=['--use_fast_math'], include_dirs=_cuda_include_dirs()) + return (cubin, ''.join(['kernel_', format(ir_obj.symbol, '')]), int(smem_bytes), int(ir_obj.threads)) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + _validate_shape(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters) + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // BLOCK_K + if not _use_splitk(dim=dim, num_n_tiles=num_n_tiles, k_tiles=k_tiles): + outputs = _splitd.launch_for_eval(inputs) + trace = dict(outputs.get('route_trace', {})) + trace['selected_route'] = _splitd.ROUTE_ID + outputs['route_trace'] = trace + return outputs + _launch_splitk(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters, num_n_tiles=num_n_tiles, k_tiles=k_tiles) + total_work = bsz * num_n_tiles * k_tiles + trace = _route_trace(inputs, total_work=total_work, grid=(min(total_work, SPLITK_GRID_CAP), 1, 1)) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': trace} + +def _validate_shape(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int) -> None: + del bsz + dtype_name = _dtype_name(inputs) + if dtype_name not in BF16_DTYPE_NAMES: + raise ValueError(''.join(['flash_kmeans_assign_highd_splitk_8de8_v1 requires bfloat16 input, got ', format(dtype_name, '')])) + if dim not in SUPPORTED_DIMS: + raise ValueError(''.join(['flash_kmeans_assign_highd_splitk_8de8_v1 supports D=', format(sorted(SUPPORTED_DIMS), ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(n_clusters, '')])) + +def _use_splitk(*, dim: int, num_n_tiles: int, k_tiles: int) -> bool: + if num_n_tiles > 16: + return False + return k_tiles >= 32 or (k_tiles >= SPLITK_MIN_K_TILES and dim >= 448) + +def _launch_splitk(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int, num_n_tiles: int, k_tiles: int) -> None: + from .._dispatch_runtime import CUDAKernel + from .._dispatch_runtime import pack_kernel_args + total_work = bsz * num_n_tiles * k_tiles + partial_scores, partial_indices = _partial_buffers(inputs, total_work) + tmap_x, tmap_c = _splitd._make_tmaps(inputs) + cubin, kernel_name, smem_bytes, threads = _compiled_partial_kernel() + partial_args = pack_kernel_args(partial_ir, x_tmap=tmap_x, c_tmap=tmap_c, c_sq=inputs['c_sq'], partial_scores=partial_scores, partial_indices=partial_indices, B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles) + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=(min(total_work, SPLITK_GRID_CAP), 1, 1), block=(threads, 1, 1), args=partial_args, shared_mem=smem_bytes) + cubin, kernel_name, smem_bytes, threads = _compiled_reduce_kernel() + reduce_args = [partial_scores, partial_indices, inputs['out'], bsz, n_points, n_clusters, num_n_tiles, k_tiles] + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=(min(bsz * num_n_tiles, SPLITK_GRID_CAP), 1, 1), block=(threads, 1, 1), args=reduce_args, shared_mem=smem_bytes) + +def _partial_buffers(inputs: dict[str, Any], total_work: int) -> tuple[Any, Any]: + import torch + cache_key = stream_cache_key(inputs, int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), int(inputs['D']), int(total_work)) + cache = inputs.setdefault('_flash_kmeans_assign_highd_splitk_8de8_v1_partials', {}) + cached = cache.get(cache_key) + if cached is not None: + return cached + scores = torch.empty((total_work, BLOCK_N), dtype=torch.float32, device=inputs['x'].device) + indices = torch.empty((total_work, BLOCK_N), dtype=torch.int32, device=inputs['x'].device) + cache[cache_key] = (scores, indices) + return (scores, indices) + +def _route_trace(inputs: dict[str, Any], *, total_work: int, grid: tuple[int, int, int]) -> dict[str, Any]: + return {'shape_key': _shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_highd_splitk_8de8_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'specialized', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_highd_splitk_8de8_v1', 'guard_condition': 'dtype == bfloat16 and D in [320, 384, 448, 512] and N % 128 == 0 and K % 256 == 0 and num_n_tiles <= 16 and (K_tiles >= 32 or (K_tiles >= 16 and D >= 448))', 'classification': 'route-ok', 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'total_tiles': total_work, 'launch_grid': grid, 'reason': 'high-D split-K seed exposes K tiles as CTA work then reduces partial argmins'} + +def compile_and_launch_flash_kmeans_assign_highd_splitk(B: int=1, N: int=512, K: int=8192, D: int=512, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(51204) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 51204}}]) + return result + +def _shape_key(inputs: dict[str, Any]) -> str: + return _splitd._shape_key(inputs) + +def _dtype_name(inputs: dict[str, Any]) -> str: + return _splitd._dtype_name(inputs) diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_splitk_blockn64_g1r4_streamdep_r63_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_splitk_blockn64_g1r4_streamdep_r63_v1.py new file mode 100644 index 00000000..25cb3e70 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_splitk_blockn64_g1r4_streamdep_r63_v1.py @@ -0,0 +1,172 @@ +"""Flash-KMeans high-D Split-K BLOCK_N=64 G1/R4 streamdep round-63 candidate. + +Minimum architecture: sm_100a. This additive child keeps the M=64 ROW_16x256B +partial geometry but reduces the grouped producer from two MMA K tiles per +work item to one. The target is to double producer work feed for no-padding +D448/D512 mid-K rows while preserving the four-lane reducer and same-stream +producer-to-reducer ordering. It is not intended for sm_120a/sm_121a where +ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +import os +from functools import lru_cache +from typing import Any +from . import flash_kmeans_assign_highd_splitk_8de8_v1 as _base +from .flash_kmeans_assign_stream import stream_cache_key +from .._dispatch_runtime import pack_kernel_args +BLOCK_N = 64 +MMA_BLOCK_K = _base.BLOCK_K +FEAT_CHUNK = _base.FEAT_CHUNK +ROW16_LOAD_COUNT = 64 +ROW16_LOGICAL_CHUNK_K = 128 +CSQ_STAGE_VEC = _base.CSQ_STAGE_VEC +X_TILE_BYTES = BLOCK_N * FEAT_CHUNK * 2 +C_TILE_BYTES = MMA_BLOCK_K * FEAT_CHUNK * 2 +CSQ_TILE_BYTES = MMA_BLOCK_K * 4 +SUPPORTED_DIMS = set(_base.SUPPORTED_DIMS) +BF16_DTYPE_NAMES = set(_base.BF16_DTYPE_NAMES) +SPLITK_GRID_CAP = _base.SPLITK_GRID_CAP +SPLITK_MIN_K_TILES = _base.SPLITK_MIN_K_TILES +SPLITK_GROUP_K_TILES = 1 +SPLITK_TILE_K = MMA_BLOCK_K * SPLITK_GROUP_K_TILES +MAX_POINT_TILES = 32 +NUM_COMPUTE_WARPS = 4 +REDUCE_LANES_PER_ROW = 4 +REDUCE_THREADS = BLOCK_N * REDUCE_LANES_PER_ROW +ROUTE_ID = 'highd_splitk_blockn64_g1r4_streamdep_r63_v1' +SEED_ID = 'highd-splitk-blockn64-g1r4-streamdep-r63-v1' +VERIFY_ENV = 'LOOM_FLASH_KMEANS_HIGHD_SPLITK_BLOCKN64_G1R4_STREAMDEP_R63_VERIFY_KERNEL' +flash_kmeans_assign_highd_splitk_partial_blockn64_g1r4_streamdep_r63_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_partial_blockn64_g1r4_streamdep_r63_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 192}')) +flash_kmeans_assign_highd_splitk_reduce_blockn64_g1r4_streamdep_r63_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_reduce_blockn64_g1r4_streamdep_r63_v1", "arg_keys": ["partial_scores", "partial_indices", "out", "B", "N", "K", "num_n_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 256}')) +partial_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_partial_blockn64_g1r4_streamdep_r63_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 192}')) +reduce_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_reduce_blockn64_g1r4_streamdep_r63_v1", "arg_keys": ["partial_scores", "partial_indices", "out", "B", "N", "K", "num_n_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 256}')) + +def _verify_export_ir() -> Any: + if os.environ.get(VERIFY_ENV) == 'reduce': + return reduce_ir + return partial_ir +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_partial_blockn64_g1r4_streamdep_r63_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 192}')) + +def _compiled_partial_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0328"}, "kernel_flash_kmeans_assign_highd_splitk_partial_blockn64_g1r4_streamdep_r63_v1", 43008, 192]}')) + +def _compiled_reduce_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0329"}, "kernel_flash_kmeans_assign_highd_splitk_reduce_blockn64_g1r4_streamdep_r63_v1", 0, 256]}')) + +@lru_cache(maxsize=1) +def _loaded_partial_kernel() -> tuple[Any, int, int]: + from .._dispatch_runtime import CUDAKernel + cubin, kernel_name, smem_bytes, threads = _compiled_partial_kernel() + return (CUDAKernel(cubin, kernel_name), smem_bytes, threads) + +@lru_cache(maxsize=1) +def _loaded_reduce_kernel() -> tuple[Any, int, int]: + from .._dispatch_runtime import CUDAKernel + cubin, kernel_name, smem_bytes, threads = _compiled_reduce_kernel() + return (CUDAKernel(cubin, kernel_name), smem_bytes, threads) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + _validate_shape(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters) + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // MMA_BLOCK_K + if not _use_blockn64_splitk(dim=dim, num_n_tiles=num_n_tiles, k_tiles=k_tiles): + outputs = _base.launch_for_eval(inputs) + trace = dict(outputs.get('route_trace', {})) + trace['selected_route'] = _base.ROUTE_ID + trace['fallback_from'] = ROUTE_ID + trace['fallback_reason'] = 'shape does not satisfy BLOCK_N=64 grouped Split-K constraints' + outputs['route_trace'] = trace + return outputs + k_slices = k_tiles // SPLITK_GROUP_K_TILES + _launch_blockn64_splitk(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters, num_n_tiles=num_n_tiles, k_tiles=k_tiles, k_slices=k_slices) + total_work = bsz * num_n_tiles * k_slices + trace = _route_trace(inputs, total_work=total_work, grid=(min(total_work, SPLITK_GRID_CAP), 1, 1), k_slices=k_slices) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': trace} + +def _validate_shape(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int) -> None: + del bsz + dtype_name = _base._dtype_name(inputs) + if dtype_name not in BF16_DTYPE_NAMES: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires bfloat16 input, got ', format(dtype_name, '')])) + if dim not in SUPPORTED_DIMS: + raise ValueError(''.join([format(ROUTE_ID, ''), ' supports D=', format(sorted(SUPPORTED_DIMS), ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % MMA_BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by MMA_BLOCK_K=', format(MMA_BLOCK_K, ''), ', got ', format(n_clusters, '')])) + +def _use_blockn64_splitk(*, dim: int, num_n_tiles: int, k_tiles: int) -> bool: + if num_n_tiles > MAX_POINT_TILES: + return False + if k_tiles % SPLITK_GROUP_K_TILES != 0: + return False + return k_tiles >= 32 or (k_tiles >= SPLITK_MIN_K_TILES and dim >= 448) + +def _launch_blockn64_splitk(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int, num_n_tiles: int, k_tiles: int, k_slices: int) -> None: + import torch + total_work = bsz * num_n_tiles * k_slices + partial_scores, partial_indices = _partial_buffers(inputs, total_work) + tmap_x, tmap_c = _make_tmaps(inputs) + stream = torch.cuda.current_stream() + partial_kernel, smem_bytes, threads = _loaded_partial_kernel() + partial_args = pack_kernel_args(partial_ir, x_tmap=tmap_x, c_tmap=tmap_c, c_sq=inputs['c_sq'], partial_scores=partial_scores, partial_indices=partial_indices, B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles, K_slices=k_slices) + partial_kernel.launch(grid=(min(total_work, SPLITK_GRID_CAP), 1, 1), block=(threads, 1, 1), args=partial_args, shared_mem=smem_bytes, stream=stream) + reduce_kernel, smem_bytes, threads = _loaded_reduce_kernel() + reduce_args = [partial_scores, partial_indices, inputs['out'], bsz, n_points, n_clusters, num_n_tiles, k_slices] + reduce_kernel.launch(grid=(min(bsz * num_n_tiles, SPLITK_GRID_CAP), 1, 1), block=(threads, 1, 1), args=reduce_args, shared_mem=smem_bytes, stream=stream) + +def _make_tmaps(inputs: dict[str, Any]) -> tuple[Any, Any]: + cache_key = (int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), int(inputs['D']), int(BLOCK_N)) + cached = inputs.get('_flash_kmeans_assign_highd_splitk_blockn64_g1r4_streamdep_r63_v1_tmaps') + if cached is not None and cached[0] == cache_key: + return (cached[1], cached[2]) + from .._dispatch_runtime import create_tensor_map_3d + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + dim = int(inputs['D']) + tmap_x = create_tensor_map_3d(inputs['x'].data_ptr(), bsz * n_points, BLOCK_N, dim, FEAT_CHUNK) + tmap_c = create_tensor_map_3d(inputs['centroids'].data_ptr(), bsz * n_clusters, MMA_BLOCK_K, dim, FEAT_CHUNK) + inputs['_flash_kmeans_assign_highd_splitk_blockn64_g1r4_streamdep_r63_v1_tmaps'] = (cache_key, tmap_x, tmap_c) + return (tmap_x, tmap_c) + +def _partial_buffers(inputs: dict[str, Any], total_work: int) -> tuple[Any, Any]: + import torch + cache_key = stream_cache_key(inputs, int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), int(inputs['D']), int(total_work), int(BLOCK_N), int(SPLITK_GROUP_K_TILES)) + cache = inputs.setdefault('_flash_kmeans_assign_highd_splitk_blockn64_g1r4_streamdep_r63_v1_partials', {}) + cached = cache.get(cache_key) + if cached is not None: + return cached + scores = torch.empty((total_work, BLOCK_N), dtype=torch.float32, device=inputs['x'].device) + indices = torch.empty((total_work, BLOCK_N), dtype=torch.int32, device=inputs['x'].device) + cache[cache_key] = (scores, indices) + return (scores, indices) + +def _route_trace(inputs: dict[str, Any], *, total_work: int, grid: tuple[int, int, int], k_slices: int) -> dict[str, Any]: + return {'shape_key': _base._shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_highd_splitk_blockn64_g1r4_streamdep_r63_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'specialized', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_highd_splitk_blockn64_g1r4_streamdep_r63_v1', 'guard_condition': 'dtype == bfloat16 and D in [320, 384, 448, 512] and N % 64 == 0 and K % 256 == 0 and num_n_tiles <= 32 and (K_tiles >= 32 or (K_tiles >= 16 and D >= 448))', 'classification': 'route-ok', 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'total_tiles': total_work, 'launch_grid': grid, 'tile_shape': {'BLOCK_N': BLOCK_N, 'BLOCK_K': SPLITK_TILE_K, 'mma_block_k': MMA_BLOCK_K, 'split_k_slices': k_slices, 'grouped_mma_k_tiles': SPLITK_GROUP_K_TILES, 'cluster_grouping': 1, 'tmem_layout': 'ROW_16x256B', 'producer_to_reducer_sync': 'same_stream_ordering', 'cuda_module_cache': 'persistent_per_process', 'producer_work_feed': '2x_g2_for_same_N_K'}, 'reason': 'BLOCK_N=64 G1/R4 tile-search variant uses ROW_16x256B score drains, single-MMA producer slices, and four-lane Split-K reducer'} + +def compile_and_launch_flash_kmeans_assign_highd_splitk_blockn64_g1r4(B: int=1, N: int=512, K: int=8192, D: int=512, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(51204) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 51204}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_splitk_blockn64_g2r4_b5a6_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_splitk_blockn64_g2r4_b5a6_v1.py new file mode 100644 index 00000000..843e93b2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_splitk_blockn64_g2r4_b5a6_v1.py @@ -0,0 +1,160 @@ +"""Flash-KMeans high-D Split-K BLOCK_N=64 G2/R4 tile-search candidate. + +Minimum architecture: sm_100a. This candidate keeps the high-D Split-K +tcgen05/TMEM score producer on the contract-visible path, keeps the M=64 +ROW_16x256B partial geometry, and tests a two-MMA grouped producer with the +four-lane Split-K reducer. The goal is to expose twice as many producer CTAs +as the G4/R4 parent without the K256 candidate's fourfold partial traffic. +It is not intended for sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +import os +from functools import lru_cache +from typing import Any +from . import flash_kmeans_assign_highd_splitk_8de8_v1 as _base +from .flash_kmeans_assign_stream import stream_cache_key +from .._dispatch_runtime import pack_kernel_args +BLOCK_N = 64 +MMA_BLOCK_K = _base.BLOCK_K +FEAT_CHUNK = _base.FEAT_CHUNK +ROW16_LOAD_COUNT = 64 +ROW16_LOGICAL_CHUNK_K = 128 +CSQ_STAGE_VEC = _base.CSQ_STAGE_VEC +X_TILE_BYTES = BLOCK_N * FEAT_CHUNK * 2 +C_TILE_BYTES = MMA_BLOCK_K * FEAT_CHUNK * 2 +CSQ_TILE_BYTES = MMA_BLOCK_K * 4 +SUPPORTED_DIMS = set(_base.SUPPORTED_DIMS) +BF16_DTYPE_NAMES = set(_base.BF16_DTYPE_NAMES) +SPLITK_GRID_CAP = _base.SPLITK_GRID_CAP +SPLITK_MIN_K_TILES = _base.SPLITK_MIN_K_TILES +SPLITK_GROUP_K_TILES = 2 +SPLITK_TILE_K = MMA_BLOCK_K * SPLITK_GROUP_K_TILES +MAX_POINT_TILES = 32 +NUM_COMPUTE_WARPS = 4 +REDUCE_LANES_PER_ROW = 4 +REDUCE_THREADS = BLOCK_N * REDUCE_LANES_PER_ROW +ROUTE_ID = 'highd_splitk_blockn64_g2r4_b5a6_v1' +SEED_ID = 'highd-splitk-blockn64-g2r4-b5a6-v1' +flash_kmeans_assign_highd_splitk_partial_blockn64_g2r4_b5a6_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_partial_blockn64_g2r4_b5a6_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 192}')) +flash_kmeans_assign_highd_splitk_reduce_blockn64_g2r4_b5a6_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_reduce_blockn64_g2r4_b5a6_v1", "arg_keys": ["partial_scores", "partial_indices", "out", "B", "N", "K", "num_n_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 256}')) +partial_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_partial_blockn64_g2r4_b5a6_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 192}')) +reduce_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_reduce_blockn64_g2r4_b5a6_v1", "arg_keys": ["partial_scores", "partial_indices", "out", "B", "N", "K", "num_n_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 256}')) + +def _verify_export_ir() -> Any: + if os.environ.get('LOOM_FLASH_KMEANS_HIGHD_SPLITK_BLOCKN64_G2R4_B5A6_VERIFY_KERNEL') == 'reduce': + return reduce_ir + return partial_ir +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_partial_blockn64_g2r4_b5a6_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 192}')) + +def _compiled_partial_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0323"}, "kernel_flash_kmeans_assign_highd_splitk_partial_blockn64_g2r4_b5a6_v1", 43008, 192]}')) + +def _compiled_reduce_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0324"}, "kernel_flash_kmeans_assign_highd_splitk_reduce_blockn64_g2r4_b5a6_v1", 0, 256]}')) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + _validate_shape(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters) + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // MMA_BLOCK_K + if not _use_blockn64_splitk(dim=dim, num_n_tiles=num_n_tiles, k_tiles=k_tiles): + outputs = _base.launch_for_eval(inputs) + trace = dict(outputs.get('route_trace', {})) + trace['selected_route'] = _base.ROUTE_ID + trace['fallback_from'] = ROUTE_ID + trace['fallback_reason'] = 'shape does not satisfy BLOCK_N=64 grouped Split-K constraints' + outputs['route_trace'] = trace + return outputs + k_slices = k_tiles // SPLITK_GROUP_K_TILES + _launch_blockn64_splitk(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters, num_n_tiles=num_n_tiles, k_tiles=k_tiles, k_slices=k_slices) + total_work = bsz * num_n_tiles * k_slices + trace = _route_trace(inputs, total_work=total_work, grid=(min(total_work, SPLITK_GRID_CAP), 1, 1), k_slices=k_slices) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': trace} + +def _validate_shape(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int) -> None: + del bsz + dtype_name = _base._dtype_name(inputs) + if dtype_name not in BF16_DTYPE_NAMES: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires bfloat16 input, got ', format(dtype_name, '')])) + if dim not in SUPPORTED_DIMS: + raise ValueError(''.join([format(ROUTE_ID, ''), ' supports D=', format(sorted(SUPPORTED_DIMS), ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % MMA_BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by MMA_BLOCK_K=', format(MMA_BLOCK_K, ''), ', got ', format(n_clusters, '')])) + +def _use_blockn64_splitk(*, dim: int, num_n_tiles: int, k_tiles: int) -> bool: + if num_n_tiles > MAX_POINT_TILES: + return False + if k_tiles % SPLITK_GROUP_K_TILES != 0: + return False + return k_tiles >= 32 or (k_tiles >= SPLITK_MIN_K_TILES and dim >= 448) + +def _launch_blockn64_splitk(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int, num_n_tiles: int, k_tiles: int, k_slices: int) -> None: + from .._dispatch_runtime import CUDAKernel + total_work = bsz * num_n_tiles * k_slices + partial_scores, partial_indices = _partial_buffers(inputs, total_work) + tmap_x, tmap_c = _make_tmaps(inputs) + cubin, kernel_name, smem_bytes, threads = _compiled_partial_kernel() + partial_args = pack_kernel_args(partial_ir, x_tmap=tmap_x, c_tmap=tmap_c, c_sq=inputs['c_sq'], partial_scores=partial_scores, partial_indices=partial_indices, B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles, K_slices=k_slices) + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=(min(total_work, SPLITK_GRID_CAP), 1, 1), block=(threads, 1, 1), args=partial_args, shared_mem=smem_bytes) + cubin, kernel_name, smem_bytes, threads = _compiled_reduce_kernel() + reduce_args = [partial_scores, partial_indices, inputs['out'], bsz, n_points, n_clusters, num_n_tiles, k_slices] + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=(min(bsz * num_n_tiles, SPLITK_GRID_CAP), 1, 1), block=(threads, 1, 1), args=reduce_args, shared_mem=smem_bytes) + +def _make_tmaps(inputs: dict[str, Any]) -> tuple[Any, Any]: + cache_key = (int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), int(inputs['D']), int(BLOCK_N)) + cached = inputs.get('_flash_kmeans_assign_highd_splitk_blockn64_g2r4_b5a6_v1_tmaps') + if cached is not None and cached[0] == cache_key: + return (cached[1], cached[2]) + from .._dispatch_runtime import create_tensor_map_3d + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + dim = int(inputs['D']) + tmap_x = create_tensor_map_3d(inputs['x'].data_ptr(), bsz * n_points, BLOCK_N, dim, FEAT_CHUNK) + tmap_c = create_tensor_map_3d(inputs['centroids'].data_ptr(), bsz * n_clusters, MMA_BLOCK_K, dim, FEAT_CHUNK) + inputs['_flash_kmeans_assign_highd_splitk_blockn64_g2r4_b5a6_v1_tmaps'] = (cache_key, tmap_x, tmap_c) + return (tmap_x, tmap_c) + +def _partial_buffers(inputs: dict[str, Any], total_work: int) -> tuple[Any, Any]: + import torch + cache_key = stream_cache_key(inputs, int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), int(inputs['D']), int(total_work), int(BLOCK_N), int(SPLITK_GROUP_K_TILES)) + cache = inputs.setdefault('_flash_kmeans_assign_highd_splitk_blockn64_g2r4_b5a6_v1_partials', {}) + cached = cache.get(cache_key) + if cached is not None: + return cached + scores = torch.empty((total_work, BLOCK_N), dtype=torch.float32, device=inputs['x'].device) + indices = torch.empty((total_work, BLOCK_N), dtype=torch.int32, device=inputs['x'].device) + cache[cache_key] = (scores, indices) + return (scores, indices) + +def _route_trace(inputs: dict[str, Any], *, total_work: int, grid: tuple[int, int, int], k_slices: int) -> dict[str, Any]: + return {'shape_key': _base._shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_highd_splitk_blockn64_g2r4_b5a6_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'specialized', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_highd_splitk_blockn64_g2r4_b5a6_v1', 'guard_condition': 'dtype == bfloat16 and D in [320, 384, 448, 512] and N % 64 == 0 and K % 256 == 0 and num_n_tiles <= 32 and K_tiles % 2 == 0 and (K_tiles >= 32 or (K_tiles >= 16 and D >= 448))', 'classification': 'route-ok', 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'total_tiles': total_work, 'launch_grid': grid, 'tile_shape': {'BLOCK_N': BLOCK_N, 'BLOCK_K': SPLITK_TILE_K, 'mma_block_k': MMA_BLOCK_K, 'split_k_slices': k_slices, 'grouped_mma_k_tiles': SPLITK_GROUP_K_TILES, 'cluster_grouping': 1, 'tmem_layout': 'ROW_16x256B'}, 'reason': 'BLOCK_N=64 G2/R4 tile-search variant uses ROW_16x256B score drains, two-MMA grouped producer slices, and four-lane Split-K reducer'} + +def compile_and_launch_flash_kmeans_assign_highd_splitk_blockn64(B: int=1, N: int=512, K: int=8192, D: int=512, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(51204) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 51204}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_splitk_blockn64_g2r4_streamdep_4f2c_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_splitk_blockn64_g2r4_streamdep_4f2c_v1.py new file mode 100644 index 00000000..5a62f2e8 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_highd_splitk_blockn64_g2r4_streamdep_4f2c_v1.py @@ -0,0 +1,104 @@ +"""Flash-KMeans high-D Split-K G2/R4 stream-dependency candidate. + +Minimum architecture: sm_100a. This candidate reuses the b5a6 BLOCK_N=64 +G2/R4 tcgen05/TMEM partial producer and four-lane Split-K reducer unchanged, +but queues the reducer after the producer on the same CUDA stream instead of +waiting on the host between launches. It also keeps the loaded CUDA modules +cached across contract iterations so the secondary launch can be queued with +less inter-kernel idle time. It is not intended for sm_120a/sm_121a where +ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +import os +from functools import lru_cache +from typing import Any +from . import flash_kmeans_assign_highd_splitk_blockn64_g2r4_b5a6_v1 as _g2r4 +from .._dispatch_runtime import pack_kernel_args +BLOCK_N = _g2r4.BLOCK_N +MMA_BLOCK_K = _g2r4.MMA_BLOCK_K +SPLITK_GROUP_K_TILES = _g2r4.SPLITK_GROUP_K_TILES +SPLITK_TILE_K = _g2r4.SPLITK_TILE_K +SPLITK_GRID_CAP = _g2r4.SPLITK_GRID_CAP +ROUTE_ID = 'highd_splitk_blockn64_g2r4_streamdep_4f2c_v1' +SEED_ID = 'highd-splitk-blockn64-g2r4-streamdep-4f2c-v1' +partial_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_partial_blockn64_g2r4_b5a6_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 192}')) +reduce_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_reduce_blockn64_g2r4_b5a6_v1", "arg_keys": ["partial_scores", "partial_indices", "out", "B", "N", "K", "num_n_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 256}')) + +def _verify_export_ir() -> Any: + if os.environ.get('LOOM_FLASH_KMEANS_HIGHD_SPLITK_BLOCKN64_G2R4_STREAMDEP_4F2C_VERIFY_KERNEL') == 'reduce': + return reduce_ir + return partial_ir +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_partial_blockn64_g2r4_b5a6_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 192}')) + +@lru_cache(maxsize=1) +def _loaded_partial_kernel() -> tuple[Any, int, int]: + from .._dispatch_runtime import CUDAKernel + cubin, kernel_name, smem_bytes, threads = _g2r4._compiled_partial_kernel() + return (CUDAKernel(cubin, kernel_name), smem_bytes, threads) + +@lru_cache(maxsize=1) +def _loaded_reduce_kernel() -> tuple[Any, int, int]: + from .._dispatch_runtime import CUDAKernel + cubin, kernel_name, smem_bytes, threads = _g2r4._compiled_reduce_kernel() + return (CUDAKernel(cubin, kernel_name), smem_bytes, threads) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + _g2r4._validate_shape(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters) + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // MMA_BLOCK_K + if not _g2r4._use_blockn64_splitk(dim=dim, num_n_tiles=num_n_tiles, k_tiles=k_tiles): + outputs = _g2r4.launch_for_eval(inputs) + trace = dict(outputs.get('route_trace', {})) + trace['selected_route'] = trace.get('selected_route', _g2r4.ROUTE_ID) + trace['fallback_from'] = ROUTE_ID + trace['fallback_reason'] = 'shape does not satisfy BLOCK_N=64 grouped Split-K constraints' + outputs['route_trace'] = trace + return outputs + k_slices = k_tiles // SPLITK_GROUP_K_TILES + total_work = bsz * num_n_tiles * k_slices + _launch_stream_ordered_splitk(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters, num_n_tiles=num_n_tiles, k_tiles=k_tiles, k_slices=k_slices, total_work=total_work) + trace = _route_trace(inputs, total_work=total_work, grid=(min(total_work, SPLITK_GRID_CAP), 1, 1), k_slices=k_slices) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': trace} + +def _launch_stream_ordered_splitk(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int, num_n_tiles: int, k_tiles: int, k_slices: int, total_work: int) -> None: + import torch + partial_scores, partial_indices = _g2r4._partial_buffers(inputs, total_work) + tmap_x, tmap_c = _g2r4._make_tmaps(inputs) + stream = torch.cuda.current_stream() + partial_kernel, smem_bytes, threads = _loaded_partial_kernel() + partial_args = pack_kernel_args(partial_ir, x_tmap=tmap_x, c_tmap=tmap_c, c_sq=inputs['c_sq'], partial_scores=partial_scores, partial_indices=partial_indices, B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles, K_slices=k_slices) + partial_kernel.launch(grid=(min(total_work, SPLITK_GRID_CAP), 1, 1), block=(threads, 1, 1), args=partial_args, shared_mem=smem_bytes, stream=stream) + reduce_kernel, smem_bytes, threads = _loaded_reduce_kernel() + reduce_args = [partial_scores, partial_indices, inputs['out'], bsz, n_points, n_clusters, num_n_tiles, k_slices] + reduce_kernel.launch(grid=(min(bsz * num_n_tiles, SPLITK_GRID_CAP), 1, 1), block=(threads, 1, 1), args=reduce_args, shared_mem=smem_bytes, stream=stream) + +def _route_trace(inputs: dict[str, Any], *, total_work: int, grid: tuple[int, int, int], k_slices: int) -> dict[str, Any]: + trace = _g2r4._route_trace(inputs, total_work=total_work, grid=grid, k_slices=k_slices) + trace.update({'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_highd_splitk_blockn64_g2r4_streamdep_4f2c_v1:launch_for_eval', 'selected_seed': SEED_ID, 'guard_id': 'guard_highd_splitk_blockn64_g2r4_streamdep_4f2c_v1', 'guard_condition': 'dtype == bfloat16 and D in [320, 384, 448, 512] and N % 64 == 0 and K % 256 == 0 and num_n_tiles <= 32 and K_tiles % 2 == 0 and (K_tiles >= 32 or (K_tiles >= 16 and D >= 448))', 'tile_shape': {'BLOCK_N': BLOCK_N, 'BLOCK_K': SPLITK_TILE_K, 'mma_block_k': MMA_BLOCK_K, 'split_k_slices': k_slices, 'grouped_mma_k_tiles': SPLITK_GROUP_K_TILES, 'cluster_grouping': 1, 'tmem_layout': 'ROW_16x256B', 'producer_to_reducer_sync': 'same_stream_ordering', 'cuda_module_cache': 'persistent_per_process'}, 'reason': 'BLOCK_N=64 G2/R4 tile-search variant reuses b5a6 IR but queues producer and reducer on one stream without an intermediate host synchronize'}) + return trace + +def compile_and_launch_flash_kmeans_assign_highd_splitk_blockn64(B: int=1, N: int=512, K: int=8192, D: int=448, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(44802) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 44802}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_lowdim_e50c_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_lowdim_e50c_v1.py new file mode 100644 index 00000000..09d29500 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_lowdim_e50c_v1.py @@ -0,0 +1,196 @@ +"""Flash-KMeans Euclidean assignment low-D seed for D=48/64/80/96/112. + +Minimum architecture: sm_100a. This candidate uses Blackwell tcgen05/TMEM +(``lm.mma``) with a Weave BF16 pack/zero-fill producer from logical +D=48/64/80/96/112 to a 128-wide MMA K tile. It is not intended for +sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +from .flash_kmeans_assign_microdim_6cd2_v1 import ir as d64_ir +from .flash_kmeans_assign_stream import stream_cache_key +BLOCK_N = 128 +BLOCK_K = 256 +SCORE_CHUNK_K = 128 +CSQ_VEC = 4 +CSQ_STAGE_VEC = 4 +FEAT_D_PAD = 128 +FEAT_D_PAD_VECS = FEAT_D_PAD // 8 +FEAT_D_DIRECT = 64 +SUPPORTED_D = {48, 64, 80, 96, 112} +NUM_COMPUTE_WARPS = 4 +X_TILE_BYTES = BLOCK_N * FEAT_D_PAD * 2 +C_TILE_BYTES = BLOCK_K * FEAT_D_PAD * 2 +CSQ_TILE_BYTES = BLOCK_K * 4 +PACK_THREADS = 256 +PACK_GRID_CAP = 4096 +_TMAP_CACHE: dict[tuple[int, int, int, int, int, int], Any] = {} +_D64_TMAP_CACHE: dict[tuple[int, int, int, int, int, int], Any] = {} +flash_kmeans_assign_lowdim_pack_e50c_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_lowdim_pack_e50c_v1", "arg_keys": ["x", "centroids", "x_pad", "c_pad", "B", "N", "D", "K", "total_x_pad", "total_c_pad"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 256}')) +flash_kmeans_assign_lowdim_e50c_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_lowdim_e50c_v1", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 100352, "constants": [], "cta_group": 1, "threads": 192}')) +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_lowdim_e50c_v1", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 100352, "constants": [], "cta_group": 1, "threads": 192}')) +pack_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_lowdim_pack_e50c_v1", "arg_keys": ["x", "centroids", "x_pad", "c_pad", "B", "N", "D", "K", "total_x_pad", "total_c_pad"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 256}')) + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as _common_cuda_include_dirs + return _common_cuda_include_dirs() + +def _create_padded_tensor_map_3d(data_ptr: int, global_height: int, shared_height: int): + import torch + from .._dispatch_runtime import create_tensor_map_3d + device_index = torch.cuda.current_device() + key = (device_index, int(data_ptr), int(global_height), int(shared_height), FEAT_D_PAD, FEAT_D_PAD) + cached = _TMAP_CACHE.get(key) + if cached is not None: + return cached + cached = create_tensor_map_3d(data_ptr, global_height, shared_height, FEAT_D_PAD, FEAT_D_PAD).to(device=torch.device('cuda', device_index)) + _TMAP_CACHE[key] = cached + return cached + +def _make_tmaps(inputs: dict[str, Any], x_pad: Any, c_pad: Any) -> tuple[Any, Any]: + cache_key = (int(x_pad.data_ptr()), int(c_pad.data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), FEAT_D_PAD) + cached = inputs.get('_flash_kmeans_assign_lowdim_e50c_v1_tmaps') + if cached is not None and cached[0] == cache_key: + return (cached[1], cached[2]) + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + tmap_x = _create_padded_tensor_map_3d(x_pad.data_ptr(), bsz * n_points, BLOCK_N) + tmap_c = _create_padded_tensor_map_3d(c_pad.data_ptr(), bsz * n_clusters, BLOCK_K) + inputs['_flash_kmeans_assign_lowdim_e50c_v1_tmaps'] = (cache_key, tmap_x, tmap_c) + return (tmap_x, tmap_c) + +def _create_d64_tensor_map_3d(data_ptr: int, global_height: int, shared_height: int): + import torch + from .._dispatch_runtime import create_tensor_map_3d + device_index = torch.cuda.current_device() + key = (device_index, int(data_ptr), int(global_height), int(shared_height), FEAT_D_DIRECT, FEAT_D_DIRECT) + cached = _D64_TMAP_CACHE.get(key) + if cached is not None: + return cached + cached = create_tensor_map_3d(data_ptr, global_height, shared_height, FEAT_D_DIRECT, FEAT_D_DIRECT).to(device=torch.device('cuda', device_index)) + _D64_TMAP_CACHE[key] = cached + return cached + +def _make_d64_tmaps(inputs: dict[str, Any]) -> tuple[Any, Any]: + cache_key = (int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), FEAT_D_DIRECT) + cached = inputs.get('_flash_kmeans_assign_lowdim_e50c_v1_d64_tmaps') + if cached is not None and cached[0] == cache_key: + return (cached[1], cached[2]) + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + tmap_x = _create_d64_tensor_map_3d(inputs['x'].data_ptr(), bsz * n_points, BLOCK_N) + tmap_c = _create_d64_tensor_map_3d(inputs['centroids'].data_ptr(), bsz * n_clusters, BLOCK_K) + inputs['_flash_kmeans_assign_lowdim_e50c_v1_d64_tmaps'] = (cache_key, tmap_x, tmap_c) + return (tmap_x, tmap_c) + +def _scratch_buffers(inputs: dict[str, Any]) -> tuple[Any, Any]: + import torch + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + cache_key = stream_cache_key(inputs, int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), bsz, n_points, n_clusters, int(inputs['D']), FEAT_D_PAD) + cache = inputs.setdefault('_flash_kmeans_assign_lowdim_e50c_v1_scratch', {}) + cached = cache.get(cache_key) + if cached is not None: + return cached + x_pad = torch.empty((bsz, n_points, FEAT_D_PAD), dtype=inputs['x'].dtype, device=inputs['x'].device) + c_pad = torch.empty((bsz, n_clusters, FEAT_D_PAD), dtype=inputs['centroids'].dtype, device=inputs['centroids'].device) + cache[cache_key] = (x_pad, c_pad) + inputs.pop('_flash_kmeans_assign_lowdim_e50c_v1_tmaps', None) + return (x_pad, c_pad) + +def _compiled_pack_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0335"}, "kernel_flash_kmeans_assign_lowdim_pack_e50c_v1", 0, 256]}')) + +def _compiled_d64_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0336"}, "kernel_flash_kmeans_assign_microdim_6cd2_v1", 51200, 192]}')) + +def _launch_d64_direct(inputs: dict[str, Any]) -> None: + from .._dispatch_runtime import CUDAKernel + from .._dispatch_runtime import pack_kernel_args + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + tmap_x, tmap_c = _make_d64_tmaps(inputs) + cubin, kernel_name, smem_bytes, threads = _compiled_d64_kernel() + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // BLOCK_K + args = pack_kernel_args(d64_ir, x_tmap=tmap_x, c_tmap=tmap_c, x_sq=inputs['x_sq'], c_sq=inputs['c_sq'], out=inputs['out'], B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles) + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=(bsz * num_n_tiles, 1, 1), block=(threads, 1, 1), args=args, shared_mem=smem_bytes) + +def _launch_pack(inputs: dict[str, Any], x_pad: Any, c_pad: Any) -> None: + from .._dispatch_runtime import CUDAKernel + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + total_x_pad = bsz * n_points * FEAT_D_PAD + total_c_pad = bsz * n_clusters * FEAT_D_PAD + work_items = max(total_x_pad, total_c_pad) + grid_x = min((work_items + PACK_THREADS - 1) // PACK_THREADS, PACK_GRID_CAP) + cubin, kernel_name, smem_bytes, threads = _compiled_pack_kernel() + args = [inputs['x'], inputs['centroids'], x_pad, c_pad, bsz, n_points, dim, n_clusters, total_x_pad, total_c_pad] + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=(grid_x, 1, 1), block=(threads, 1, 1), args=args, shared_mem=smem_bytes) + +def _compiled_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0337"}, "kernel_flash_kmeans_assign_lowdim_e50c_v1", 100352, 192]}')) + +def launch_for_eval(inputs: dict[str, Any]) -> Any: + from .._dispatch_runtime import CUDAKernel + from .._dispatch_runtime import pack_kernel_args + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + dtype = str(inputs.get('dtype', getattr(inputs['x'], 'dtype', 'bfloat16'))).replace('torch.', '') + if dtype not in {'bfloat16', 'bf16'}: + raise ValueError(''.join(['flash_kmeans_assign_lowdim_e50c_v1 requires bfloat16 input, got ', format(dtype, '')])) + if dim not in SUPPORTED_D: + raise ValueError(''.join(['flash_kmeans_assign_lowdim_e50c_v1 requires D in ', format(sorted(SUPPORTED_D), ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(n_clusters, '')])) + if dim == FEAT_D_DIRECT: + _launch_d64_direct(inputs) + return {'cluster_ids': inputs['out']} + x_pad, c_pad = _scratch_buffers(inputs) + _launch_pack(inputs, x_pad, c_pad) + tmap_x, tmap_c = _make_tmaps(inputs, x_pad, c_pad) + cubin, kernel_name, smem_bytes, threads = _compiled_kernel() + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // BLOCK_K + grid = (bsz * num_n_tiles, 1, 1) + block = (threads, 1, 1) + args = pack_kernel_args(ir, x_tmap=tmap_x, c_tmap=tmap_c, x_sq=inputs['x_sq'], c_sq=inputs['c_sq'], out=inputs['out'], B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles) + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=grid, block=block, args=args, shared_mem=smem_bytes) + return {'cluster_ids': inputs['out']} + +def compile_and_launch_flash_kmeans_assign_lowdim(B: int=4, N: int=1024, K: int=512, D: int=64, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(6401) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 6401}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_microdim_6cd2_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_microdim_6cd2_v1.py new file mode 100644 index 00000000..f83d77b2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_microdim_6cd2_v1.py @@ -0,0 +1,147 @@ +"""Flash-KMeans Euclidean assignment micro-D seed for D=16/32. + +Minimum architecture: sm_100a. This candidate uses Blackwell tcgen05/TMEM +(``lm.mma``) with a Weave BF16 pack/zero-fill producer from logical D=16/32 +to a single 64-wide MMA K slice. It is not intended for sm_120a/sm_121a where +ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +from .flash_kmeans_assign_stream import stream_cache_key +BLOCK_N = 128 +BLOCK_K = 256 +SCORE_CHUNK_K = 128 +CSQ_VEC = 4 +CSQ_STAGE_VEC = 4 +FEAT_D_PAD = 64 +FEAT_D_PAD_VECS = FEAT_D_PAD // 8 +SUPPORTED_D = {16, 32} +NUM_COMPUTE_WARPS = 4 +X_TILE_BYTES = BLOCK_N * FEAT_D_PAD * 2 +C_TILE_BYTES = BLOCK_K * FEAT_D_PAD * 2 +CSQ_TILE_BYTES = BLOCK_K * 4 +PACK_THREADS = 256 +PACK_GRID_CAP = 4096 +_TMAP_CACHE: dict[tuple[int, int, int, int, int, int], Any] = {} +flash_kmeans_assign_microdim_pack_6cd2_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_microdim_pack_6cd2_v1", "arg_keys": ["x", "centroids", "x_pad", "c_pad", "B", "N", "D", "K", "total_x_pad", "total_c_pad"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 256}')) +flash_kmeans_assign_microdim_6cd2_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_microdim_6cd2_v1", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 51200, "constants": [], "cta_group": 1, "threads": 192}')) +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_microdim_6cd2_v1", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 51200, "constants": [], "cta_group": 1, "threads": 192}')) +pack_ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_microdim_pack_6cd2_v1", "arg_keys": ["x", "centroids", "x_pad", "c_pad", "B", "N", "D", "K", "total_x_pad", "total_c_pad"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 0, "constants": [], "cta_group": 1, "threads": 256}')) + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as _common_cuda_include_dirs + return _common_cuda_include_dirs() + +def _create_padded_tensor_map_3d(data_ptr: int, global_height: int, shared_height: int): + import torch + from .._dispatch_runtime import create_tensor_map_3d + device_index = torch.cuda.current_device() + key = (device_index, int(data_ptr), int(global_height), int(shared_height), FEAT_D_PAD, FEAT_D_PAD) + cached = _TMAP_CACHE.get(key) + if cached is not None: + return cached + cached = create_tensor_map_3d(data_ptr, global_height, shared_height, FEAT_D_PAD, FEAT_D_PAD).to(device=torch.device('cuda', device_index)) + _TMAP_CACHE[key] = cached + return cached + +def _make_tmaps(inputs: dict[str, Any], x_pad: Any, c_pad: Any) -> tuple[Any, Any]: + cache_key = (int(x_pad.data_ptr()), int(c_pad.data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), FEAT_D_PAD) + cached = inputs.get('_flash_kmeans_assign_microdim_6cd2_v1_tmaps') + if cached is not None and cached[0] == cache_key: + return (cached[1], cached[2]) + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + tmap_x = _create_padded_tensor_map_3d(x_pad.data_ptr(), bsz * n_points, BLOCK_N) + tmap_c = _create_padded_tensor_map_3d(c_pad.data_ptr(), bsz * n_clusters, BLOCK_K) + inputs['_flash_kmeans_assign_microdim_6cd2_v1_tmaps'] = (cache_key, tmap_x, tmap_c) + return (tmap_x, tmap_c) + +def _scratch_buffers(inputs: dict[str, Any]) -> tuple[Any, Any]: + import torch + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + cache_key = stream_cache_key(inputs, int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), bsz, n_points, n_clusters, int(inputs['D'])) + cache = inputs.setdefault('_flash_kmeans_assign_microdim_6cd2_v1_scratch', {}) + cached = cache.get(cache_key) + if cached is not None: + return cached + x_pad = torch.empty((bsz, n_points, FEAT_D_PAD), dtype=inputs['x'].dtype, device=inputs['x'].device) + c_pad = torch.empty((bsz, n_clusters, FEAT_D_PAD), dtype=inputs['centroids'].dtype, device=inputs['centroids'].device) + cache[cache_key] = (x_pad, c_pad) + inputs.pop('_flash_kmeans_assign_microdim_6cd2_v1_tmaps', None) + return (x_pad, c_pad) + +def _compiled_pack_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0333"}, "kernel_flash_kmeans_assign_microdim_pack_6cd2_v1", 0, 256]}')) + +def _launch_pack(inputs: dict[str, Any], x_pad: Any, c_pad: Any) -> None: + from .._dispatch_runtime import CUDAKernel + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + total_x_pad = bsz * n_points * FEAT_D_PAD + total_c_pad = bsz * n_clusters * FEAT_D_PAD + work_items = max(total_x_pad, total_c_pad) + grid_x = min((work_items + PACK_THREADS - 1) // PACK_THREADS, PACK_GRID_CAP) + cubin, kernel_name, smem_bytes, threads = _compiled_pack_kernel() + args = [inputs['x'], inputs['centroids'], x_pad, c_pad, bsz, n_points, dim, n_clusters, total_x_pad, total_c_pad] + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=(grid_x, 1, 1), block=(threads, 1, 1), args=args, shared_mem=smem_bytes) + +def _compiled_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0334"}, "kernel_flash_kmeans_assign_microdim_6cd2_v1", 51200, 192]}')) + +def launch_for_eval(inputs: dict[str, Any]) -> Any: + from .._dispatch_runtime import CUDAKernel + from .._dispatch_runtime import pack_kernel_args + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + dtype = str(inputs.get('dtype', getattr(inputs['x'], 'dtype', 'bfloat16'))).replace('torch.', '') + if dtype not in {'bfloat16', 'bf16'}: + raise ValueError(''.join(['flash_kmeans_assign_microdim_6cd2_v1 requires bfloat16 input, got ', format(dtype, '')])) + if dim not in SUPPORTED_D: + raise ValueError(''.join(['flash_kmeans_assign_microdim_6cd2_v1 requires D in ', format(sorted(SUPPORTED_D), ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(n_clusters, '')])) + x_pad, c_pad = _scratch_buffers(inputs) + _launch_pack(inputs, x_pad, c_pad) + tmap_x, tmap_c = _make_tmaps(inputs, x_pad, c_pad) + cubin, kernel_name, smem_bytes, threads = _compiled_kernel() + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // BLOCK_K + grid = (bsz * num_n_tiles, 1, 1) + block = (threads, 1, 1) + args = pack_kernel_args(ir, x_tmap=tmap_x, c_tmap=tmap_c, x_sq=inputs['x_sq'], c_sq=inputs['c_sq'], out=inputs['out'], B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles) + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=grid, block=block, args=args, shared_mem=smem_bytes) + return {'cluster_ids': inputs['out']} + +def compile_and_launch_flash_kmeans_assign_microdim(B: int=4, N: int=1024, K: int=512, D: int=16, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(1601) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 1601}}]) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_microdim_direct_9c0d_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_microdim_direct_9c0d_v1.py new file mode 100644 index 00000000..4aa5adfc --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_microdim_direct_9c0d_v1.py @@ -0,0 +1,101 @@ +"""Flash-KMeans Euclidean assignment direct micro-D seed for D=16/32. + +Minimum architecture: sm_100a. This candidate uses Blackwell tcgen05/TMEM +(``lm.mma``) for D=16/32 by staging logical BF16 rows directly into a +64-wide padded K-major shared-memory tile. It is not intended for +sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +BLOCK_N = 128 +BLOCK_K = 256 +SCORE_CHUNK_K = 128 +CSQ_VEC = 4 +CSQ_STAGE_VEC = 4 +FEAT_D_PAD = 64 +FEAT_D_PAD_VECS = FEAT_D_PAD // 8 +STAGE_VEC = 8 +SUPPORTED_D = {16, 32} +NUM_COMPUTE_WARPS = 4 +X_TILE_BYTES = BLOCK_N * FEAT_D_PAD * 2 +C_TILE_BYTES = BLOCK_K * FEAT_D_PAD * 2 +CSQ_TILE_BYTES = BLOCK_K * 4 +ROUTE_ID = 'microdim_direct_staged_9c0d_v1' +SEED_ID = 'microdim-direct-staged-9c0d-v1' +BF16_DTYPE_NAMES = {'bfloat16', 'bf16', 'torch.bfloat16'} +flash_kmeans_assign_microdim_direct_9c0d_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_microdim_direct_9c0d_v1", "arg_keys": ["x", "centroids", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 51200, "constants": [], "cta_group": 1, "threads": 192}')) +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_microdim_direct_9c0d_v1", "arg_keys": ["x", "centroids", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 51200, "constants": [], "cta_group": 1, "threads": 192}')) + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as _common_cuda_include_dirs + return _common_cuda_include_dirs() + +def _compiled_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0332"}, "kernel_flash_kmeans_assign_microdim_direct_9c0d_v1", 51200, 192]}')) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + from .._dispatch_runtime import CUDAKernel + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + dtype_name = _dtype_name(inputs) + if dtype_name not in BF16_DTYPE_NAMES: + raise ValueError(''.join(['flash_kmeans_assign_microdim_direct_9c0d_v1 requires bfloat16 input, got ', format(dtype_name, '')])) + if dim not in SUPPORTED_D: + raise ValueError(''.join(['flash_kmeans_assign_microdim_direct_9c0d_v1 requires D in ', format(sorted(SUPPORTED_D), ''), ', got ', format(dim, '')])) + if n_points % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(n_points, '')])) + if n_clusters % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(n_clusters, '')])) + cubin, kernel_name, smem_bytes, threads = _compiled_kernel() + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // BLOCK_K + grid = (bsz * num_n_tiles, 1, 1) + block = (threads, 1, 1) + args = [inputs['x'], inputs['centroids'], inputs['x_sq'], inputs['c_sq'], inputs['out'], bsz, n_points, dim, n_clusters, num_n_tiles, k_tiles] + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=grid, block=block, args=args, shared_mem=smem_bytes) + trace = _route_trace(inputs, total_tiles=bsz * num_n_tiles, grid=grid) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': trace} + +def compile_and_launch_flash_kmeans_assign_microdim_direct(B: int=4, N: int=1024, K: int=512, D: int=16, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + torch.manual_seed(1601) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=[{'label': ''.join(['manual_b', format(B, ''), '_n', format(N, ''), '_k', format(K, ''), '_d', format(D, '')]), 'params': {'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'seed': 1601}}]) + return result + +def _route_trace(inputs: dict[str, Any], *, total_tiles: int, grid: tuple[int, int, int]) -> dict[str, Any]: + return {'shape_key': _shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_microdim_direct_9c0d_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'specialized', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_microdim_direct_staged_9c0d_v1', 'guard_condition': 'dtype == bfloat16 and D in [16, 32] and N % 128 == 0 and K % 256 == 0', 'classification': 'route-ok', 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'total_tiles': total_tiles, 'launch_grid': grid, 'reason': 'micro-D direct SMEM staging feeds one padded tcgen05 score tile without global scratch packing'} + +def _shape_key(inputs: dict[str, Any]) -> str: + label = inputs.get('label') + if label: + return str(label) + return ''.join(['b', format(int(inputs['B']), ''), '_n', format(int(inputs['N']), ''), '_k', format(int(inputs['K']), ''), '_d', format(int(inputs['D']), '')]) + +def _dtype_name(inputs: dict[str, Any]) -> str: + dtype = inputs.get('dtype') + if dtype is not None: + return str(dtype).replace('torch.', '') + x = inputs.get('x') + if x is not None and hasattr(x, 'dtype'): + return str(x.dtype).replace('torch.', '') + return 'bfloat16' diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_microdim_hybrid_9c0d_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_microdim_hybrid_9c0d_v1.py new file mode 100644 index 00000000..3b279b95 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_microdim_hybrid_9c0d_v1.py @@ -0,0 +1,93 @@ +"""Flash-KMeans Euclidean assignment guarded micro-D seed for D=16/32. + +Minimum architecture: sm_100a. This bucket seed combines two Weave-only +tcgen05/TMEM paths: direct SMEM staging for small K=512 rows where global +scratch packing is overhead-dominant, and the existing 6cd2 pack+TMA path for +large/high-K rows where centroid reuse makes TMA staging faster. It is not +intended for sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from typing import Any +from . import flash_kmeans_assign_microdim_6cd2_v1 as _pack +from . import flash_kmeans_assign_microdim_direct_9c0d_v1 as _direct +BLOCK_N = _direct.BLOCK_N +BLOCK_K = _direct.BLOCK_K +SUPPORTED_D = _direct.SUPPORTED_D +BF16_DTYPE_NAMES = _direct.BF16_DTYPE_NAMES +ROUTE_ID = 'microdim_hybrid_9c0d_v1' +SEED_ID = 'microdim-hybrid-9c0d-v1' +DIRECT_CHILD_ROUTE_ID = _direct.ROUTE_ID +PACK_CHILD_ROUTE_ID = 'microdim_pack_6cd2_v1' +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_microdim_direct_9c0d_v1", "arg_keys": ["x", "centroids", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 51200, "constants": [], "cta_group": 1, "threads": 192}')) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz, n_points, dim, n_clusters, dtype_name = _shape_fields(inputs) + _validate_supported_shape(N=n_points, D=dim, K=n_clusters, dtype=dtype_name) + if _use_direct_path(N=n_points, K=n_clusters): + child_outputs = _direct.launch_for_eval(inputs) + child_route = DIRECT_CHILD_ROUTE_ID + child_seed = _direct.SEED_ID + else: + child_outputs = _pack.launch_for_eval(inputs) + child_route = PACK_CHILD_ROUTE_ID + child_seed = 'microdim-pack-6cd2-v1' + normalized = _normalize_outputs(child_outputs, inputs) + trace = _route_trace(inputs, child_route=child_route, child_seed=child_seed, direct_path=child_route == DIRECT_CHILD_ROUTE_ID, total_tiles=bsz * (n_points // BLOCK_N)) + normalized['selected_route'] = ROUTE_ID + normalized['route_trace'] = trace + return normalized + +def _use_direct_path(*, N: int, K: int) -> bool: + return K == 512 and N <= 2432 + +def _validate_supported_shape(*, N: int, D: int, K: int, dtype: Any) -> None: + dtype_name = str(dtype).replace('torch.', '') + if dtype_name not in BF16_DTYPE_NAMES: + raise ValueError(''.join(['flash_kmeans_assign_microdim_hybrid_9c0d_v1 requires bfloat16 input, got ', format(dtype, '')])) + if D not in SUPPORTED_D: + raise ValueError(''.join(['flash_kmeans_assign_microdim_hybrid_9c0d_v1 requires D in ', format(sorted(SUPPORTED_D), ''), ', got ', format(D, '')])) + if N % BLOCK_N != 0: + raise ValueError(''.join(['N must be divisible by BLOCK_N=', format(BLOCK_N, ''), ', got ', format(N, '')])) + if K % BLOCK_K != 0: + raise ValueError(''.join(['K must be divisible by BLOCK_K=', format(BLOCK_K, ''), ', got ', format(K, '')])) + +def _normalize_outputs(outputs: Any, inputs: dict[str, Any]) -> dict[str, Any]: + if outputs is None: + return {'cluster_ids': inputs['out']} + if hasattr(outputs, 'shape'): + return {'cluster_ids': outputs} + if isinstance(outputs, dict): + normalized = dict(outputs) + if 'cluster_ids' not in normalized and 'out' in normalized: + normalized['cluster_ids'] = normalized['out'] + if 'cluster_ids' in normalized: + return normalized + raise TypeError("flash_kmeans_assign microdim hybrid route must return cluster_ids or write inputs['out']") + +def _route_trace(inputs: dict[str, Any], *, child_route: str, child_seed: str, direct_path: bool, total_tiles: int) -> dict[str, Any]: + return {'shape_key': _shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_microdim_hybrid_9c0d_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'specialized', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_microdim_hybrid_9c0d_v1', 'guard_condition': 'dtype == bfloat16 and D in [16, 32] and N % 128 == 0 and K % 256 == 0', 'classification': 'route-ok', 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'child_route': child_route, 'child_seed': child_seed, 'direct_staging': direct_path, 'total_tiles': total_tiles, 'reason': 'K=512 short rows use direct SMEM staging; larger rows retain the 6cd2 pack+TMA path for centroid reuse'} + +def _shape_fields(inputs: dict[str, Any]) -> tuple[int, int, int, int, str]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + dtype_name = _dtype_name(inputs) + return (bsz, n_points, dim, n_clusters, dtype_name) + +def _shape_key(inputs: dict[str, Any]) -> str: + label = inputs.get('label') + if label: + return str(label) + return ''.join(['b', format(int(inputs['B']), ''), '_n', format(int(inputs['N']), ''), '_k', format(int(inputs['K']), ''), '_d', format(int(inputs['D']), '')]) + +def _dtype_name(inputs: dict[str, Any]) -> str: + dtype = inputs.get('dtype') + if dtype is not None: + return str(dtype).replace('torch.', '') + x = inputs.get('x') + if x is not None and hasattr(x, 'dtype'): + return str(x.dtype).replace('torch.', '') + return 'bfloat16' diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_microdim_pipeline4_08f9_v4.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_microdim_pipeline4_08f9_v4.py new file mode 100644 index 00000000..334fcc9d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_microdim_pipeline4_08f9_v4.py @@ -0,0 +1,155 @@ +"""Flash-KMeans high-N D16 four-stage pipeline with D32 raw-TMA sibling. + +Minimum architecture: sm_100a. D16 keeps one CTA per 128-point tile and +uses exact K16, 32B-swizzled TMA operands. Four centroid stages and four TMEM +score stages overlap the eight K64 centroid chunks, while three point stages +feed a bounded persistent grid. A single cooperative K512 c_sq preload +replaces the per-stage load/synchronize sequence, and four independent score +chains shorten the per-row argmax dependency. D32 retains the correct +raw-TMA v1 child. There is no global scratch, partial output, or +host/reference sidecar. This module +is not intended for sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +from . import flash_kmeans_assign_microdim_raw_tma_08f9_v1 as _d32 +BLOCK_N = 128 +BLOCK_K = 64 +FEAT_D = 16 +NUM_C_STAGES = 4 +NUM_ACC_STAGES = 4 +NUM_X_STAGES = 3 +NUM_COMPUTE_WARPS = 4 +THREADS = 192 +GRID_CAP = 152 +X_STAGE_BYTES = BLOCK_N * FEAT_D * 2 +C_STAGE_BYTES = BLOCK_K * FEAT_D * 2 +X_TOTAL_BYTES = X_STAGE_BYTES * NUM_X_STAGES +C_TOTAL_BYTES = C_STAGE_BYTES * NUM_C_STAGES +CSQ_ELEMS = 512 +CSQ_BYTES = CSQ_ELEMS * 4 +ROUTE_ID = 'microdim_pipeline4_08f9_v4' +SEED_ID = 'microdim-pipeline4-08f9-v4' +D16_CHILD_ROUTE_ID = 'microdim_d16_pipeline4_k64_08f9_v4' +D32_CHILD_ROUTE_ID = _d32.ROUTE_ID +TARGET_SHAPE = 'post_d895_d16_b8_n65536_k512_d16' +TARGET_SHAPES = (TARGET_SHAPE, 'post_d895_d32_b8_n65536_k512_d32') +BF16_DTYPE_NAMES = {'bfloat16', 'bf16', 'torch.bfloat16'} +_TMAP_CACHE: dict[tuple[int, int, int, int, int, int], Any] = {} +flash_kmeans_assign_microdim_d16_pipeline4_08f9_v4 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_microdim_d16_pipeline4_08f9_v4", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 23552, "constants": [], "cta_group": 1, "threads": 192}')) +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_microdim_d16_pipeline4_08f9_v4", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 23552, "constants": [], "cta_group": 1, "threads": 192}')) + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as _common_cuda_include_dirs + return _common_cuda_include_dirs() + +def _create_d16_tmap(data_ptr: int, *, global_height: int, shared_height: int): + import torch + from .._dispatch_runtime import create_tensor_map_3d_32b + device_index = torch.cuda.current_device() + key = (device_index, int(data_ptr), int(global_height), int(shared_height), FEAT_D, FEAT_D) + cached = _TMAP_CACHE.get(key) + if cached is not None: + return cached + cached = create_tensor_map_3d_32b(data_ptr, global_height, shared_height, FEAT_D, FEAT_D).to(device=torch.device('cuda', device_index)) + _TMAP_CACHE[key] = cached + return cached + +def _make_d16_tmaps(inputs: dict[str, Any]) -> tuple[Any, Any]: + cache_key = (int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['K']), FEAT_D) + cached = inputs.get('_flash_kmeans_assign_microdim_pipeline4_08f9_v4_tmaps') + if cached is not None and cached[0] == cache_key: + return (cached[1], cached[2]) + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + tmap_x = _create_d16_tmap(inputs['x'].data_ptr(), global_height=bsz * n_points, shared_height=BLOCK_N) + tmap_c = _create_d16_tmap(inputs['centroids'].data_ptr(), global_height=bsz * n_clusters, shared_height=BLOCK_K) + inputs['_flash_kmeans_assign_microdim_pipeline4_08f9_v4_tmaps'] = (cache_key, tmap_x, tmap_c) + return (tmap_x, tmap_c) + +def _compiled_d16_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0330"}, "kernel_flash_kmeans_assign_microdim_d16_pipeline4_08f9_v4", 23552, 192]}')) + +def _launch_d16(inputs: dict[str, Any]) -> None: + from .._dispatch_runtime import CUDAKernel + from .._dispatch_runtime import pack_kernel_args + bsz = int(inputs['B']) + n_points = int(inputs['N']) + n_clusters = int(inputs['K']) + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // BLOCK_K + tmap_x, tmap_c = _make_d16_tmaps(inputs) + cubin, kernel_name, smem_bytes, threads = _compiled_d16_kernel() + args = pack_kernel_args(ir, x_tmap=tmap_x, c_tmap=tmap_c, c_sq=inputs['c_sq'], out=inputs['out'], B=bsz, N=n_points, D=FEAT_D, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles) + grid = (min(bsz * num_n_tiles, GRID_CAP), 1, 1) + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=grid, block=(threads, 1, 1), args=args, shared_mem=smem_bytes) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + _validate_supported_shape(B=bsz, N=n_points, D=dim, K=n_clusters, dtype=_dtype_name(inputs)) + if dim == FEAT_D: + _launch_d16(inputs) + child_route = D16_CHILD_ROUTE_ID + child_seed = 'microdim-d16-pipeline4-k64-08f9-v4' + else: + child_outputs = _d32.launch_for_eval(inputs) + if child_outputs['cluster_ids'].data_ptr() != inputs['out'].data_ptr(): + raise AssertionError('D32 raw-TMA child must write the caller-owned output') + child_route = D32_CHILD_ROUTE_ID + child_seed = _d32.SEED_ID + trace = _route_trace(inputs, child_route=child_route, child_seed=child_seed, pipelined=dim == FEAT_D) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': trace} + +def compile_and_launch_flash_kmeans_assign_microdim_b8_highn_2(*, benchmark: bool=False) -> dict[str, Any]: + from .._dispatch_runtime import evaluate + shapes = _exact_shapes(benchmark=benchmark) + + def checked(inputs: dict[str, Any]) -> dict[str, Any]: + outputs = launch_for_eval(inputs) + if outputs['cluster_ids'].data_ptr() != inputs['out'].data_ptr(): + raise AssertionError('microdim pipeline4 seed must write the caller-owned output') + if outputs['selected_route'] != ROUTE_ID: + raise AssertionError('microdim pipeline4 seed returned unexpected route metadata') + if outputs['route_trace']['shape_key'] != inputs['label']: + raise AssertionError('microdim pipeline4 seed returned the wrong shape key') + return outputs + payload = evaluate(checked, shapes=shapes, correctness=True, benchmark=benchmark, time_triton_baseline=False) + return {'passed': bool(payload['correctness']['all_correct']), 'contract_eval': payload} + +def _exact_shapes(*, benchmark: bool) -> list[dict[str, Any]]: + return [{'label': TARGET_SHAPES[0], 'params': {'B': 8, 'N': 65536, 'K': 512, 'D': 16, 'dtype': 'bfloat16', 'seed': 21601, 'check_correctness': True, 'benchmark': benchmark}}, {'label': TARGET_SHAPES[1], 'params': {'B': 8, 'N': 65536, 'K': 512, 'D': 32, 'dtype': 'bfloat16', 'seed': 23201, 'check_correctness': True, 'benchmark': benchmark}}] + +def _validate_supported_shape(*, B: int, N: int, D: int, K: int, dtype: Any) -> None: + dtype_name = str(dtype).replace('torch.', '') + if dtype_name not in BF16_DTYPE_NAMES: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires bfloat16 input, got ', format(dtype, '')])) + if B != 8 or N != 65536 or K != 512 or (D not in {16, 32}): + raise ValueError(''.join([format(ROUTE_ID, ''), ' owns B=8, N=65536, K=512, D in [16, 32]'])) + if N % BLOCK_N != 0 or K % BLOCK_K != 0: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires N % ', format(BLOCK_N, ''), ' == 0 and K % ', format(BLOCK_K, ''), ' == 0'])) + +def _route_trace(inputs: dict[str, Any], *, child_route: str, child_seed: str, pipelined: bool) -> dict[str, Any]: + return {'shape_key': _shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_microdim_pipeline4_08f9_v4:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'specialized', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_microdim_pipeline4_08f9_v4', 'guard_condition': 'B == 8 and N == 65536 and K == 512 and D in [16, 32] and dtype == bfloat16', 'classification': 'route-ok', 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'child_route': child_route, 'child_seed': child_seed, 'pipelined': pipelined, 'caller_owned_output': True, 'residual_contract_regions': [], 'reason': 'D16 preloads K512 c_sq once and overlaps exact K16 K64 tiles through four C/TMEM stages; D32 retains the one-CTA raw-TMA sibling'} + +def _shape_key(inputs: dict[str, Any]) -> str: + label = inputs.get('label') + if label: + return str(label) + return ''.join(['b', format(int(inputs['B']), ''), '_n', format(int(inputs['N']), ''), '_k', format(int(inputs['K']), ''), '_d', format(int(inputs['D']), '')]) + +def _dtype_name(inputs: dict[str, Any]) -> str: + dtype = inputs.get('dtype') + if dtype is not None: + return str(dtype).replace('torch.', '') + x = inputs.get('x') + if x is not None and hasattr(x, 'dtype'): + return str(x.dtype).replace('torch.', '') + return 'bfloat16' diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_microdim_raw_tma_08f9_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_microdim_raw_tma_08f9_v1.py new file mode 100644 index 00000000..b1136878 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_microdim_raw_tma_08f9_v1.py @@ -0,0 +1,114 @@ +"""Flash-KMeans high-N micro-D raw-TMA seed for D=16/32. + +Minimum architecture: sm_100a. This candidate keeps the proven 6cd2 +tcgen05/TMEM score and argmax schedule, but replaces its global pad kernel and +rank-3 padded tensor maps with rank-2 tensor maps over the raw logical rows. +The 64-wide TMA box zero-fills columns beyond logical D, so the contract path +is one assignment kernel with no global scratch handoff. It is not intended +for sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from functools import lru_cache +from typing import Any +BLOCK_N = 128 +BLOCK_K = 256 +SCORE_CHUNK_K = 128 +CSQ_VEC = 4 +CSQ_STAGE_VEC = 4 +FEAT_D_PAD = 64 +SUPPORTED_D = {16, 32} +NUM_COMPUTE_WARPS = 4 +X_TILE_BYTES = BLOCK_N * FEAT_D_PAD * 2 +C_TILE_BYTES = BLOCK_K * FEAT_D_PAD * 2 +CSQ_ELEMS = 512 +CSQ_TILE_BYTES = CSQ_ELEMS * 4 +ROUTE_ID = 'microdim_raw_tma_08f9_v1' +SEED_ID = 'microdim-raw-tma-08f9-v1' +BF16_DTYPE_NAMES = {'bfloat16', 'bf16', 'torch.bfloat16'} +TARGET_SHAPE = 'post_d895_d16_b8_n65536_k512_d16' +TARGET_SHAPES = (TARGET_SHAPE, 'post_d895_d32_b8_n65536_k512_d32') +flash_kmeans_assign_microdim_raw_tma_08f9_v1 = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_microdim_raw_tma_08f9_v1", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 52224, "constants": [], "cta_group": 1, "threads": 192}')) +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_microdim_raw_tma_08f9_v1", "arg_keys": ["x_tmap", "c_tmap", "x_sq", "c_sq", "out", "B", "N", "D", "K", "num_n_tiles", "K_tiles"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 52224, "constants": [], "cta_group": 1, "threads": 192}')) +_TMAP_CACHE: dict[tuple[int, int, int, int, int, int], Any] = {} + +def _cuda_include_dirs() -> list[str]: + from .._dispatch_runtime import _cuda_include_dirs as _common_cuda_include_dirs + return _common_cuda_include_dirs() + +def _create_raw_tensor_map(data_ptr: int, *, global_height: int, shared_height: int, dim: int): + import torch + from .._dispatch_runtime import create_tensor_map + device_index = torch.cuda.current_device() + key = (device_index, int(data_ptr), int(global_height), int(shared_height), int(dim), FEAT_D_PAD) + cached = _TMAP_CACHE.get(key) + if cached is not None: + return cached + cached = create_tensor_map(data_ptr, dim, global_height, FEAT_D_PAD, shared_height, dim * 2).to(device=torch.device('cuda', device_index)) + _TMAP_CACHE[key] = cached + return cached + +def _make_tmaps(inputs: dict[str, Any]) -> tuple[Any, Any]: + cache_key = (int(inputs['x'].data_ptr()), int(inputs['centroids'].data_ptr()), int(inputs['B']), int(inputs['N']), int(inputs['D']), int(inputs['K'])) + cached = inputs.get('_flash_kmeans_assign_microdim_raw_tma_08f9_v1_tmaps') + if cached is not None and cached[0] == cache_key: + return (cached[1], cached[2]) + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + tmap_x = _create_raw_tensor_map(inputs['x'].data_ptr(), global_height=bsz * n_points, shared_height=BLOCK_N, dim=dim) + tmap_c = _create_raw_tensor_map(inputs['centroids'].data_ptr(), global_height=bsz * n_clusters, shared_height=BLOCK_K, dim=dim) + inputs['_flash_kmeans_assign_microdim_raw_tma_08f9_v1_tmaps'] = (cache_key, tmap_x, tmap_c) + return (tmap_x, tmap_c) + +def _compiled_kernel() -> tuple[bytes, str, int, int]: + return _decode_capture(_json_loads('{"__tuple__": [{"__kernel_source__": "dispatch_kernel_0331"}, "kernel_flash_kmeans_assign_microdim_raw_tma_08f9_v1", 52224, 192]}')) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + from .._dispatch_runtime import CUDAKernel + from .._dispatch_runtime import pack_kernel_args + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + dtype_name = _dtype_name(inputs) + _validate_supported_shape(B=bsz, N=n_points, D=dim, K=n_clusters, dtype=dtype_name) + tmap_x, tmap_c = _make_tmaps(inputs) + cubin, kernel_name, smem_bytes, threads = _compiled_kernel() + num_n_tiles = n_points // BLOCK_N + k_tiles = n_clusters // BLOCK_K + grid = (bsz * num_n_tiles, 1, 1) + args = pack_kernel_args(ir, x_tmap=tmap_x, c_tmap=tmap_c, x_sq=inputs['x_sq'], c_sq=inputs['c_sq'], out=inputs['out'], B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles) + with CUDAKernel(cubin, kernel_name) as kernel: + kernel.launch(grid=grid, block=(threads, 1, 1), args=args, shared_mem=smem_bytes) + trace = _route_trace(inputs, total_tiles=bsz * num_n_tiles, grid=grid) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': trace} + +def _validate_supported_shape(*, B: int, N: int, D: int, K: int, dtype: Any) -> None: + dtype_name = str(dtype).replace('torch.', '') + if dtype_name not in BF16_DTYPE_NAMES: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires bfloat16 input, got ', format(dtype, '')])) + if B != 8 or N != 65536 or K != 512 or (D not in SUPPORTED_D): + raise ValueError(''.join([format(ROUTE_ID, ''), ' owns B=8, N=65536, K=512, D in ', format(sorted(SUPPORTED_D), '')])) + if N % BLOCK_N != 0 or K % BLOCK_K != 0: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires N % ', format(BLOCK_N, ''), ' == 0 and K % ', format(BLOCK_K, ''), ' == 0'])) + +def _route_trace(inputs: dict[str, Any], *, total_tiles: int, grid: tuple[int, int, int]) -> dict[str, Any]: + return {'shape_key': _shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_microdim_raw_tma_08f9_v1:launch_for_eval', 'selected_seed': SEED_ID, 'expected_seed': None, 'route_kind': 'specialized', 'route_source': 'shape-specific-seed', 'guard_id': 'guard_microdim_raw_tma_08f9_v1', 'guard_condition': 'B == 8 and N == 65536 and K == 512 and D in [16, 32] and dtype == bfloat16', 'classification': 'route-ok', 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'raw_tma_zero_fill': True, 'global_padding_scratch': False, 'total_tiles': total_tiles, 'launch_grid': grid, 'reason': 'raw rank-2 TMA zero-fills the 64-wide micro-D tile and removes the pack-kernel handoff'} + +def _shape_key(inputs: dict[str, Any]) -> str: + label = inputs.get('label') + if label: + return str(label) + return ''.join(['b', format(int(inputs['B']), ''), '_n', format(int(inputs['N']), ''), '_k', format(int(inputs['K']), ''), '_d', format(int(inputs['D']), '')]) + +def _dtype_name(inputs: dict[str, Any]) -> str: + dtype = inputs.get('dtype') + if dtype is not None: + return str(dtype).replace('torch.', '') + x = inputs.get('x') + if x is not None and hasattr(x, 'dtype'): + return str(x.dtype).replace('torch.', '') + return 'bfloat16' diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_no_padding_portfolio_r51_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_no_padding_portfolio_r51_v1.py new file mode 100644 index 00000000..63d1268a --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_no_padding_portfolio_r51_v1.py @@ -0,0 +1,98 @@ +"""Flash-KMeans no-padding high-D Split-K portfolio round-51 candidate. + +Minimum architecture: sm_100a. This additive bucket-kernel composition keeps +Blackwell tcgen05/TMEM children on the contract-visible path for the no-padding +high-D Split-K rows from the round-50 handoff plus expanded no-padding heldout +rows. It routes D448 paired to the round-47 dual-TMEM X-reuse producer plus +round-39 R1 reducer, D512 paired to the recorded round-39 incumbent, and the +other high-D rows directly through the 4f2c streamdep child. It is not intended +for sm_120a/sm_121a where ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +import os +from typing import Any +from . import flash_kmeans_assign_highd_paired_ownerreduce_r39_v1 as _r39 +from . import flash_kmeans_assign_highd_paired_xreuse_dualtmem_r47_v1 as _r47 +from . import flash_kmeans_assign_highd_splitk_8de8_v1 as _base +from . import flash_kmeans_assign_highd_splitk_blockn64_g2r4_streamdep_4f2c_v1 as _streamdep +ROUTE_ID = 'flash_kmeans_assign_no_padding_portfolio_r51_v1' +SEED_ID = 'no-padding-highd-portfolio-r51-v1' +VERIFY_ENV = 'LOOM_FLASH_KMEANS_NO_PADDING_PORTFOLIO_R51_VERIFY_KERNEL' +NO_PADDING_HIGHD_SHAPES = {(1, 512, 8192, 320): 'd320_hugek_b1_n512_k8192_d320', (1, 1024, 8192, 320): 'd320_hugek_b1_n1024_k8192_d320', (1, 2048, 8192, 320): 'd320_hugek_b1_n2048_k8192_d320', (1, 512, 8192, 384): 'd384_hugek_b1_n512_k8192_d384', (1, 768, 8192, 384): 'd384_hugek_b1_n768_k8192_d384', (1, 1024, 8192, 384): 'd384_hugek_b1_n1024_k8192_d384', (1, 512, 4096, 448): 'd448_midk_b1_n512_k4096_d448', (1, 1024, 4096, 448): 'd448_midk_b1_n1024_k4096_d448', (1, 512, 8192, 448): 'd448_hugek_b1_n512_k8192_d448', (1, 1024, 8192, 448): 'd448_hugek_b1_n1024_k8192_d448', (1, 2048, 4096, 448): 'd448_paired_b1_n2048_k4096_d448', (1, 512, 4096, 512): 'd512_midk_b1_n512_k4096_d512', (1, 1024, 4096, 512): 'd512_midk_b1_n1024_k4096_d512', (1, 512, 8192, 512): 'd512_hugek_b1_n512_k8192_d512', (1, 1024, 8192, 512): 'd512_hugek_b1_n1024_k8192_d512', (1, 2048, 4096, 512): 'd512_paired_b1_n2048_k4096_d512'} +NO_PADDING_HIGHD_EVAL_SHAPES = [{'label': 'd320_hugek_b1_n512_k8192_d320', 'params': {'B': 1, 'N': 512, 'D': 320, 'K': 8192, 'dtype': 'bfloat16', 'seed': 32004}}, {'label': 'd320_hugek_b1_n1024_k8192_d320', 'params': {'B': 1, 'N': 1024, 'D': 320, 'K': 8192, 'dtype': 'bfloat16', 'seed': 32005}}, {'label': 'd320_hugek_b1_n2048_k8192_d320', 'params': {'B': 1, 'N': 2048, 'D': 320, 'K': 8192, 'dtype': 'bfloat16', 'seed': 32006}}, {'label': 'd384_hugek_b1_n512_k8192_d384', 'params': {'B': 1, 'N': 512, 'D': 384, 'K': 8192, 'dtype': 'bfloat16', 'seed': 38403}}, {'label': 'd384_hugek_b1_n768_k8192_d384', 'params': {'B': 1, 'N': 768, 'D': 384, 'K': 8192, 'dtype': 'bfloat16', 'seed': 38404}}, {'label': 'd384_hugek_b1_n1024_k8192_d384', 'params': {'B': 1, 'N': 1024, 'D': 384, 'K': 8192, 'dtype': 'bfloat16', 'seed': 38405}}, {'label': 'd448_midk_b1_n512_k4096_d448', 'params': {'B': 1, 'N': 512, 'D': 448, 'K': 4096, 'dtype': 'bfloat16', 'seed': 44803}}, {'label': 'd448_midk_b1_n1024_k4096_d448', 'params': {'B': 1, 'N': 1024, 'D': 448, 'K': 4096, 'dtype': 'bfloat16', 'seed': 44805}}, {'label': 'd448_hugek_b1_n512_k8192_d448', 'params': {'B': 1, 'N': 512, 'D': 448, 'K': 8192, 'dtype': 'bfloat16', 'seed': 44802}}, {'label': 'd448_hugek_b1_n1024_k8192_d448', 'params': {'B': 1, 'N': 1024, 'D': 448, 'K': 8192, 'dtype': 'bfloat16', 'seed': 44806}}, {'label': 'd448_paired_b1_n2048_k4096_d448', 'params': {'B': 1, 'N': 2048, 'D': 448, 'K': 4096, 'dtype': 'bfloat16', 'seed': 44801}}, {'label': 'd512_midk_b1_n512_k4096_d512', 'params': {'B': 1, 'N': 512, 'D': 512, 'K': 4096, 'dtype': 'bfloat16', 'seed': 51203}}, {'label': 'd512_midk_b1_n1024_k4096_d512', 'params': {'B': 1, 'N': 1024, 'D': 512, 'K': 4096, 'dtype': 'bfloat16', 'seed': 51205}}, {'label': 'd512_hugek_b1_n512_k8192_d512', 'params': {'B': 1, 'N': 512, 'D': 512, 'K': 8192, 'dtype': 'bfloat16', 'seed': 51204}}, {'label': 'd512_hugek_b1_n1024_k8192_d512', 'params': {'B': 1, 'N': 1024, 'D': 512, 'K': 8192, 'dtype': 'bfloat16', 'seed': 51206}}, {'label': 'd512_paired_b1_n2048_k4096_d512', 'params': {'B': 1, 'N': 2048, 'D': 512, 'K': 4096, 'dtype': 'bfloat16', 'seed': 51202}}] + +def _verify_export_ir() -> Any: + verify_kernel = os.environ.get(VERIFY_ENV) + if verify_kernel == 'd448_reduce': + return _r47.reduce_ir + if verify_kernel == 'd512_partial': + return _r39._r2.partial_ir + if verify_kernel == 'd512_reduce': + return _r39._r2.reduce_ir + if verify_kernel == 'hugek_partial': + return _streamdep.partial_ir + if verify_kernel == 'hugek_reduce': + return _streamdep.reduce_ir + return _r47.partial_ir +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_xreuse_dualtmem_producer_r47_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_keys", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 75776, "constants": [], "cta_group": 1, "threads": 192}')) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + _validate_shape(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters) + if _is_d448_paired(bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters): + outputs = _r47.launch_for_eval(inputs) + trace = _wrap_child_trace(inputs, child_outputs=outputs, route_kind='shape-specific-seed-composition', selected_child='r47_dualtmem_xreuse_plus_r39_reduce1', reason='D448 paired row uses the round-47 dual-TMEM X-reuse producer and the round-39 unrolled R1 reducer') + return {'cluster_ids': outputs['cluster_ids'], 'selected_route': ROUTE_ID, 'route_trace': trace} + if _is_d512_paired(bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters): + outputs = _r39.launch_for_eval(inputs) + trace = _wrap_child_trace(inputs, child_outputs=outputs, route_kind='shape-specific-seed-composition', selected_child='r39_ownerreduce_incumbent', reason='D512 paired row keeps the recorded round-39 incumbent route') + return {'cluster_ids': outputs['cluster_ids'], 'selected_route': ROUTE_ID, 'route_trace': trace} + outputs = _streamdep.launch_for_eval(inputs) + trace = _wrap_child_trace(inputs, child_outputs=outputs, route_kind='shape-specific-seed-composition', selected_child='4f2c_streamdep_child', reason='huge-K rows use the same-process fastest direct stream-dependency high-D Split-K child') + return {'cluster_ids': outputs['cluster_ids'], 'selected_route': ROUTE_ID, 'route_trace': trace} + +def _validate_shape(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int) -> None: + dtype_name = _base._dtype_name(inputs) + if dtype_name not in _base.BF16_DTYPE_NAMES: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires bfloat16 input, got ', format(dtype_name, '')])) + if (bsz, n_points, n_clusters, dim) not in NO_PADDING_HIGHD_SHAPES: + raise ValueError(''.join([format(ROUTE_ID, ''), ' is exact-shape only for the no-padding high-D bucket, got B=', format(bsz, ''), ', N=', format(n_points, ''), ', K=', format(n_clusters, ''), ', D=', format(dim, '')])) + +def _is_d448_paired(*, bsz: int, n_points: int, dim: int, n_clusters: int) -> bool: + return bsz == 1 and n_points == 2048 and (n_clusters == 4096) and (dim == 448) + +def _is_d512_paired(*, bsz: int, n_points: int, dim: int, n_clusters: int) -> bool: + return bsz == 1 and n_points == 2048 and (n_clusters == 4096) and (dim == 512) + +def _wrap_child_trace(inputs: dict[str, Any], *, child_outputs: dict[str, Any], route_kind: str, selected_child: str, reason: str) -> dict[str, Any]: + child_trace = dict(child_outputs.get('route_trace', {})) + child_route = child_outputs.get('selected_route') or child_trace.get('selected_route') + trace = dict(child_trace) + trace.update({'shape_key': _base._shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_no_padding_portfolio_r51_v1:launch_for_eval', 'selected_seed': SEED_ID, 'route_kind': route_kind, 'route_source': 'shape-specific-seed-portfolio', 'guard_id': 'guard_no_padding_highd_splitk_portfolio_r51_v1', 'guard_condition': 'exact no-padding high-D bucket: D320/D384 hugeK, D448/D512 midK/hugeK/paired, all N/K aligned without padding', 'classification': 'route-ok', 'child_route': child_route, 'child_entrypoint': child_trace.get('selected_entrypoint'), 'child_guard_id': child_trace.get('guard_id'), 'child_tile_shape': child_trace.get('tile_shape'), 'portfolio_selected_child': selected_child, 'dispatcher_kernel_ms': None, 'shape_specific_kernel_ms': None, 'relative_speedup_vs_baseline': None, 'reason': reason}) + return trace + +def compile_and_launch_flash_kmeans_assign_no_padding_portfolio_r51(B: int=1, N: int=2048, K: int=4096, D: int=448, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + seed = {(1, 512, 8192, 320): 32004, (1, 1024, 8192, 320): 32005, (1, 2048, 8192, 320): 32006, (1, 512, 8192, 384): 38403, (1, 768, 8192, 384): 38404, (1, 1024, 8192, 384): 38405, (1, 512, 4096, 448): 44803, (1, 1024, 4096, 448): 44805, (1, 512, 8192, 448): 44802, (1, 1024, 8192, 448): 44806, (1, 2048, 4096, 448): 44801, (1, 512, 4096, 512): 51203, (1, 1024, 4096, 512): 51205, (1, 512, 8192, 512): 51204, (1, 1024, 8192, 512): 51206, (1, 2048, 4096, 512): 51202}[B, N, K, D] + torch.manual_seed(seed) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': NO_PADDING_HIGHD_SHAPES[B, N, K, D], 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=NO_PADDING_HIGHD_EVAL_SHAPES, correctness=True, benchmark=True, flashlib_baseline=True, benchmark_warmup_ms=100, benchmark_ms=1000, time_triton_baseline=False) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_no_padding_portfolio_r52_d448_gridcap160_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_no_padding_portfolio_r52_d448_gridcap160_v1.py new file mode 100644 index 00000000..4c412191 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_no_padding_portfolio_r52_d448_gridcap160_v1.py @@ -0,0 +1,109 @@ +"""Flash-KMeans no-padding high-D Split-K portfolio round-52 candidate. + +Minimum architecture: sm_100a. This additive bucket-kernel variant keeps the +round-51 no-padding high-D denominator but changes the D448 paired route to +launch the round-47 dual-TMEM X-reuse producer with a 160-CTA grid cap before +the round-39 R1 packed-key reducer. D512 paired and non-D448 rows delegate to +the round-51 portfolio children. It is not intended for sm_120a/sm_121a where +ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +import os +from typing import Any +from . import flash_kmeans_assign_highd_paired_xreuse_dualtmem_r47_v1 as _r47 +from . import flash_kmeans_assign_highd_splitk_blockn64_g2r4_b5a6_v1 as _g2r4 +from . import flash_kmeans_assign_no_padding_portfolio_r51_v1 as _r51 +from .._dispatch_runtime import pack_kernel_args +ROUTE_ID = 'flash_kmeans_assign_no_padding_portfolio_r52_d448_gridcap160_v1' +SEED_ID = 'no-padding-highd-portfolio-r52-d448-gridcap160-v1' +VERIFY_ENV = 'LOOM_FLASH_KMEANS_NO_PADDING_PORTFOLIO_R52_D448_GRIDCAP160_VERIFY_KERNEL' +PRODUCER_GRID_CAP_D448 = 160 +NO_PADDING_HIGHD_SHAPES = _r51.NO_PADDING_HIGHD_SHAPES +NO_PADDING_HIGHD_EVAL_SHAPES = _r51.NO_PADDING_HIGHD_EVAL_SHAPES + +def _verify_export_ir() -> Any: + verify_kernel = os.environ.get(VERIFY_ENV) + if verify_kernel == 'd448_reduce': + return _r47.reduce_ir + if verify_kernel == 'd512_partial': + return _r47._r39._r2.partial_ir + if verify_kernel == 'd512_reduce': + return _r47._r39._r2.reduce_ir + if verify_kernel == 'hugek_partial': + return _r51._streamdep.partial_ir + if verify_kernel == 'hugek_reduce': + return _r51._streamdep.reduce_ir + return _r47.partial_ir +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_paired_xreuse_dualtmem_producer_r47_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_keys", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 75776, "constants": [], "cta_group": 1, "threads": 192}')) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + _r51._validate_shape(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters) + if _r51._is_d448_paired(bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters): + num_n_tiles = n_points // _r47.BLOCK_N + k_tiles = n_clusters // _r47.MMA_BLOCK_K + k_slices = k_tiles // _r47.SPLITK_GROUP_K_TILES + if k_slices != _r47.PAIRED_K_SLICES: + raise ValueError(''.join([format(ROUTE_ID, ''), ' requires K_slices=', format(_r47.PAIRED_K_SLICES, ''), ', got ', format(k_slices, '')])) + total_work = bsz * num_n_tiles * k_slices + producer_grid, reducer_grid = _launch_d448_gridcap160(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters, num_n_tiles=num_n_tiles, k_tiles=k_tiles, k_slices=k_slices, total_work=total_work) + trace = _route_trace_d448_gridcap160(inputs, total_work=total_work, producer_grid=producer_grid, reducer_grid=reducer_grid, k_slices=k_slices) + return {'cluster_ids': inputs['out'], 'selected_route': ROUTE_ID, 'route_trace': trace} + outputs = _r51.launch_for_eval(inputs) + trace = _wrap_r51_child_trace(inputs, child_outputs=outputs, reason='non-D448 rows delegate to the round-51 no-padding portfolio') + return {'cluster_ids': outputs['cluster_ids'], 'selected_route': ROUTE_ID, 'route_trace': trace} + +def _launch_d448_gridcap160(inputs: dict[str, Any], *, bsz: int, n_points: int, dim: int, n_clusters: int, num_n_tiles: int, k_tiles: int, k_slices: int, total_work: int) -> tuple[tuple[int, int, int], tuple[int, int, int]]: + import torch + partial_keys = _r47._r1._partial_key_buffer(inputs, total_work) + tmap_x, tmap_c = _g2r4._make_tmaps(inputs) + stream = torch.cuda.current_stream() + partial_kernel, smem_bytes, threads = _r47._loaded_partial_key_kernel() + partial_args = pack_kernel_args(_r47.partial_ir, x_tmap=tmap_x, c_tmap=tmap_c, c_sq=inputs['c_sq'], partial_keys=partial_keys, B=bsz, N=n_points, D=dim, K=n_clusters, num_n_tiles=num_n_tiles, K_tiles=k_tiles, K_slices=k_slices) + producer_grid = (min(total_work, PRODUCER_GRID_CAP_D448), 1, 1) + partial_kernel.launch(grid=producer_grid, block=(threads, 1, 1), args=partial_args, shared_mem=smem_bytes, stream=stream) + reduce_kernel, smem_bytes, threads = _r47._r39._loaded_reduce1_kernel() + reduce_args = [partial_keys, inputs['out'], bsz, n_points, n_clusters, num_n_tiles, k_slices] + reducer_grid = (min(bsz * num_n_tiles, _r47.SPLITK_GRID_CAP), 1, 1) + reduce_kernel.launch(grid=reducer_grid, block=(threads, 1, 1), args=reduce_args, shared_mem=smem_bytes, stream=stream) + return (producer_grid, reducer_grid) + +def _route_trace_d448_gridcap160(inputs: dict[str, Any], *, total_work: int, producer_grid: tuple[int, int, int], reducer_grid: tuple[int, int, int], k_slices: int) -> dict[str, Any]: + trace = _r47._route_trace_d448(inputs, total_work=total_work, producer_grid=producer_grid, reducer_grid=reducer_grid, k_slices=k_slices) + trace.update({'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_no_padding_portfolio_r52_d448_gridcap160_v1:launch_for_eval', 'selected_seed': SEED_ID, 'route_source': 'shape-specific-seed-portfolio', 'guard_id': 'guard_no_padding_highd_splitk_portfolio_r52_d448_gridcap160_v1', 'guard_condition': 'exact no-padding high-D bucket; D448 paired uses R47 producer gridcap160, other rows delegate to round-51', 'portfolio_selected_child': 'r47_dualtmem_xreuse_gridcap160_plus_r39_reduce1', 'reason': 'D448 paired row keeps the round-47 dual-TMEM X-reuse producer/reducer ABI but raises the producer launch cap from 128 to 160 CTAs to reduce persistent work underfeed'}) + tile_shape = dict(trace.get('tile_shape', {})) + tile_shape.update({'producer_grid_cap': PRODUCER_GRID_CAP_D448, 'producer_grid_work_items_per_cta_floor': total_work // PRODUCER_GRID_CAP_D448, 'producer_grid_remainder_work_items': total_work % PRODUCER_GRID_CAP_D448}) + trace['tile_shape'] = tile_shape + return trace + +def _wrap_r51_child_trace(inputs: dict[str, Any], *, child_outputs: dict[str, Any], reason: str) -> dict[str, Any]: + child_trace = dict(child_outputs.get('route_trace', {})) + trace = dict(child_trace) + trace.update({'shape_key': _r51._base._shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_no_padding_portfolio_r52_d448_gridcap160_v1:launch_for_eval', 'selected_seed': SEED_ID, 'route_source': 'shape-specific-seed-portfolio', 'guard_id': 'guard_no_padding_highd_splitk_portfolio_r52_d448_gridcap160_v1', 'portfolio_selected_child': child_trace.get('portfolio_selected_child'), 'child_route': child_outputs.get('selected_route') or child_trace.get('selected_route'), 'child_entrypoint': child_trace.get('selected_entrypoint'), 'child_guard_id': child_trace.get('guard_id'), 'reason': reason}) + return trace + +def compile_and_launch_flash_kmeans_assign_no_padding_portfolio_r52_d448_gridcap160(B: int=1, N: int=2048, K: int=4096, D: int=448, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + seed = {(1, 768, 8192, 384): 38404, (1, 512, 8192, 448): 44802, (1, 2048, 4096, 448): 44801, (1, 512, 8192, 512): 51204, (1, 2048, 4096, 512): 51202}[B, N, K, D] + torch.manual_seed(seed) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': NO_PADDING_HIGHD_SHAPES[B, N, K, D], 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=NO_PADDING_HIGHD_EVAL_SHAPES, correctness=True, benchmark=True, flashlib_baseline=True, benchmark_warmup_ms=100, benchmark_ms=1000, time_triton_baseline=False) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1.py new file mode 100644 index 00000000..6c2a1c8a --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1.py @@ -0,0 +1,73 @@ +"""Flash-KMeans no-padding high-D round-63 G1/R4 D512 N512 portfolio. + +Minimum architecture: sm_100a. This additive bucket-kernel variant keeps the +round-52 no-padding portfolio for established rows, but routes only the +D512 K=4096 N512 mid-K row through the round-63 G1/R4 streamdep child to +increase producer work feed. It is not intended for sm_120a/sm_121a where +ptxas rejects tcgen05 instructions. +""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +import os +from typing import Any +from . import flash_kmeans_assign_highd_splitk_blockn64_g1r4_streamdep_r63_v1 as _g1 +from . import flash_kmeans_assign_no_padding_portfolio_r52_d448_gridcap160_v1 as _r52 +ROUTE_ID = 'flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1' +SEED_ID = 'no-padding-highd-portfolio-r63-g1-d512n512-v1' +VERIFY_ENV = 'LOOM_FLASH_KMEANS_NO_PADDING_PORTFOLIO_R63_G1_D512N512_VERIFY_KERNEL' +NO_PADDING_HIGHD_SHAPES = _r52.NO_PADDING_HIGHD_SHAPES +NO_PADDING_HIGHD_EVAL_SHAPES = _r52.NO_PADDING_HIGHD_EVAL_SHAPES + +def _verify_export_ir() -> Any: + verify_kernel = os.environ.get(VERIFY_ENV) + if verify_kernel == 'g1_reduce': + return _g1.reduce_ir + if verify_kernel == 'r52': + return _r52.ir + return _g1.partial_ir +ir = _decode_capture(_json_loads('{"__ir__": "flash_kmeans_assign_highd_splitk_partial_blockn64_g1r4_streamdep_r63_v1", "arg_keys": ["x_tmap", "c_tmap", "c_sq", "partial_scores", "partial_indices", "B", "N", "D", "K", "num_n_tiles", "K_tiles", "K_slices"], "cluster_dims": [1, 1, 1], "computed_smem_bytes": 43008, "constants": [], "cta_group": 1, "threads": 192}')) + +def launch_for_eval(inputs: dict[str, Any]) -> dict[str, Any]: + bsz = int(inputs['B']) + n_points = int(inputs['N']) + dim = int(inputs['D']) + n_clusters = int(inputs['K']) + _r52._r51._validate_shape(inputs, bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters) + if _is_g1_d512n512_target(bsz=bsz, n_points=n_points, dim=dim, n_clusters=n_clusters): + outputs = _g1.launch_for_eval(inputs) + trace = _wrap_child_trace(inputs, child_outputs=outputs, selected_child='r63_blockn64_g1r4_streamdep_child', reason='D512 N512 mid-K no-padding row uses the round-63 G1/R4 streamdep child to double producer work feed versus G2/R4') + return {'cluster_ids': outputs['cluster_ids'], 'selected_route': ROUTE_ID, 'route_trace': trace} + outputs = _r52.launch_for_eval(inputs) + trace = _wrap_child_trace(inputs, child_outputs=outputs, selected_child=outputs.get('route_trace', {}).get('portfolio_selected_child', 'r52_portfolio_child'), reason='non-target rows delegate unchanged to the round-52 no-padding portfolio') + return {'cluster_ids': outputs['cluster_ids'], 'selected_route': ROUTE_ID, 'route_trace': trace} + +def _is_g1_d512n512_target(*, bsz: int, n_points: int, dim: int, n_clusters: int) -> bool: + return bsz == 1 and n_points == 512 and (dim == 512) and (n_clusters == 4096) + +def _wrap_child_trace(inputs: dict[str, Any], *, child_outputs: dict[str, Any], selected_child: str, reason: str) -> dict[str, Any]: + child_trace = dict(child_outputs.get('route_trace', {})) + trace = dict(child_trace) + trace.update({'shape_key': _r52._r51._base._shape_key(inputs), 'selected_route': ROUTE_ID, 'selected_entrypoint': 'loom.examples.weave.flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1:launch_for_eval', 'selected_seed': SEED_ID, 'route_source': 'shape-specific-seed-portfolio', 'guard_id': 'guard_no_padding_highd_splitk_portfolio_r63_g1_d512n512_v1', 'guard_condition': 'exact no-padding high-D bucket; only D512 K4096 N512 uses G1/R4 streamdep child, other rows delegate to R52', 'portfolio_selected_child': selected_child, 'child_route': child_outputs.get('selected_route') or child_trace.get('selected_route'), 'child_entrypoint': child_trace.get('selected_entrypoint'), 'child_guard_id': child_trace.get('guard_id'), 'child_tile_shape': child_trace.get('tile_shape'), 'reason': reason}) + return trace + +def compile_and_launch_flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512(B: int=1, N: int=512, K: int=4096, D: int=512, *, benchmark: bool=False) -> dict[str, Any]: + import torch + if not torch.cuda.is_available(): + raise RuntimeError('CUDA GPU required') + seed = {(1, 512, 8192, 320): 32004, (1, 1024, 8192, 320): 32005, (1, 2048, 8192, 320): 32006, (1, 512, 8192, 384): 38403, (1, 768, 8192, 384): 38404, (1, 1024, 8192, 384): 38405, (1, 512, 4096, 448): 44803, (1, 1024, 4096, 448): 44805, (1, 512, 8192, 448): 44802, (1, 1024, 8192, 448): 44806, (1, 2048, 4096, 448): 44801, (1, 512, 4096, 512): 51203, (1, 1024, 4096, 512): 51205, (1, 512, 8192, 512): 51204, (1, 1024, 8192, 512): 51206, (1, 2048, 4096, 512): 51202}[B, N, K, D] + torch.manual_seed(seed) + x = torch.randn((B, N, D), dtype=torch.bfloat16, device='cuda').contiguous() + centroids = torch.randn((B, K, D), dtype=torch.bfloat16, device='cuda').contiguous() + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + out = torch.empty((B, N), dtype=torch.int32, device='cuda') + inputs = {'label': NO_PADDING_HIGHD_SHAPES[B, N, K, D], 'B': B, 'N': N, 'D': D, 'K': K, 'dtype': 'bfloat16', 'x': x, 'centroids': centroids, 'x_sq': x_sq, 'c_sq': c_sq, 'out': out} + launch_for_eval(inputs) + ref_dist = x_sq.unsqueeze(-1) + c_sq.unsqueeze(1) - 2.0 * torch.einsum('bnd,bkd->bnk', x.float(), centroids.float()) + ref = ref_dist.clamp_min(0.0).argmin(dim=-1).to(torch.int32) + result: dict[str, Any] = {'passed': bool(torch.equal(out, ref))} + if benchmark: + from .._dispatch_runtime import evaluate + result['contract_eval'] = evaluate(launch_for_eval, shapes=NO_PADDING_HIGHD_EVAL_SHAPES, correctness=True, benchmark=True, flashlib_baseline=True, benchmark_warmup_ms=100, benchmark_ms=1000, time_triton_baseline=False) + return result diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_stream.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_stream.py new file mode 100644 index 00000000..e390f055 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch/flash_kmeans_assign_stream.py @@ -0,0 +1,40 @@ +"""Current-stream ownership helpers for Flash-KMeans production routes.""" +from __future__ import annotations +from json import loads as _json_loads +from .._dispatch_runtime import _capture_cuTensorMapEncodeTiled, _decode_capture, _import_dispatch_module, _ir_proxy +from collections.abc import Iterator +from contextlib import contextmanager +from typing import Any + +def current_stream(inputs: dict[str, Any]) -> Any: + """Resolve the caller's current stream on the input tensor device.""" + import torch + return torch.cuda.current_stream(device=inputs['x'].device) + +def stream_cache_key(inputs: dict[str, Any], *parts: int) -> tuple[int, ...]: + """Prefix a cache key with the exact CUDA device and stream handle.""" + import torch + prepared_key = inputs.get('_flash_kmeans_assign_prepared_stream_key') + if prepared_key is not None: + device_index, stream_handle = prepared_key + else: + device = inputs['x'].device + device_index = getattr(device, 'index', None) + if device_index is None: + device_index = torch.cuda.current_device() + stream_handle = current_stream(inputs).cuda_stream + return (int(device_index), int(stream_handle), *(int(part) for part in parts)) + +@contextmanager +def bind_current_stream(inputs: dict[str, Any]) -> Iterator[Any]: + """Keep a production child pipeline on the caller's entry stream.""" + import torch + device = inputs['x'].device + stream = current_stream(inputs) + device_index = getattr(device, 'index', None) + current_device = torch.cuda.current_device() + if device_index is None or int(device_index) == int(current_device): + yield stream + return + with torch.cuda.device(device), torch.cuda.stream(stream): + yield stream diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch_runtime.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch_runtime.py new file mode 100644 index 00000000..c57622d1 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch_runtime.py @@ -0,0 +1,1151 @@ +from __future__ import annotations +_KERNEL_ALIAS_BY_IR_NAME = {'flash_kmeans_assign_lowdim_pack_e50c_v1': 'dispatch_kernel_0000', 'flash_kmeans_assign_lowdim_e50c_v1': 'dispatch_kernel_0001', 'flash_kmeans_assign_cleanroom_tcgen05_v10': 'dispatch_kernel_0002', 'flash_kmeans_assign_d160_pad192_pack_f9b2_v1': 'dispatch_kernel_0004', 'flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1': 'dispatch_kernel_0005', 'flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1': 'dispatch_kernel_0008', 'flash_kmeans_assign_highd_splitd_6fcf_v1': 'dispatch_kernel_0009', 'flash_kmeans_assign_highd_splitk_partial_blockn64_g2r4_b5a6_v1': 'dispatch_kernel_0011', 'flash_kmeans_assign_highd_splitk_reduce_blockn64_g2r4_b5a6_v1': 'dispatch_kernel_0012', 'flash_kmeans_assign_microdim_pack_6cd2_v1': 'dispatch_kernel_0013', 'flash_kmeans_assign_microdim_6cd2_v1': 'dispatch_kernel_0014', 'flash_kmeans_assign_gap_pad_pack_v1': 'dispatch_kernel_0015', 'flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1': 'dispatch_kernel_0016', 'flash_kmeans_assign_cleanroom_tcgen05_v15': 'dispatch_kernel_0048', 'flash_kmeans_assign_highd_splitk_partial_blockn64_g1r4_streamdep_r63_v1': 'dispatch_kernel_0059', 'flash_kmeans_assign_highd_splitk_reduce_blockn64_g1r4_streamdep_r63_v1': 'dispatch_kernel_0060', 'flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1': 'dispatch_kernel_0101', 'flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1': 'dispatch_kernel_0109', 'flash_kmeans_assign_highd_paired_xreuse_dualtmem_producer_r47_v1': 'dispatch_kernel_0118', 'flash_kmeans_assign_highd_paired_ownerreduce_r39_reduce1_unroll_v1': 'dispatch_kernel_0119', 'flash_kmeans_assign_highd_paired_packedpartial_producer_7b3c_v1': 'dispatch_kernel_0121', 'flash_kmeans_assign_highd_paired_packedpartial_reduce_r2_7b3c_v1': 'dispatch_kernel_0122', 'flash_kmeans_assign_microdim_d16_pipeline4_08f9_v4': 'dispatch_kernel_0126', 'flash_kmeans_assign_microdim_raw_tma_08f9_v1': 'dispatch_kernel_0129', 'flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1': 'dispatch_kernel_0156', 'flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4': 'dispatch_kernel_0229', 'flash_kmeans_assign_d288_exactd_a532_v1': 'dispatch_kernel_0236', 'flash_kmeans_assign_d288_splitk_cta_0438_v1_partial': 'dispatch_kernel_0237', 'flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce': 'dispatch_kernel_0238', 'flash_kmeans_assign_d480_splitk_partial_d32k256_v1': 'dispatch_kernel_0258', 'flash_kmeans_assign_d480_splitk_reduce_d32k256_v1': 'dispatch_kernel_0259', 'flash_kmeans_assign_highd_splitk_partial_8de8_v1': 'dispatch_kernel_0277', 'flash_kmeans_assign_highd_splitk_reduce_8de8_v1': 'dispatch_kernel_0278', 'flash_kmeans_assign_microdim_direct_9c0d_v1': 'dispatch_kernel_0332', 'flash_kmeans_assign_d416_exactd_splitd_a4a579d1_v2': 'dispatch_kernel_0347', 'flash_kmeans_assign_d112_k1024_owner_local_warp_mma_distinct_workfeed_c829_v1': 'dispatch_kernel_0353', 'flash_kmeans_assign_d112_k512_two_owner_peer_only_mma_f826_v1': 'dispatch_kernel_0354', 'flash_kmeans_assign_d112_shared_point_dual_issuer_tcgen05_6e1e_v1': 'dispatch_kernel_0355'} +_KERNEL_ALIAS_BY_REQUEST = {'{"computed_smem_bytes":0,"constants":[],"ir_name":"flash_kmeans_assign_lowdim_pack_e50c_v1","kwargs":{"smem_bytes":0,"validate":false},"threads":256}': 'dispatch_kernel_0000', '{"computed_smem_bytes":100352,"constants":[],"ir_name":"flash_kmeans_assign_lowdim_e50c_v1","kwargs":{"smem_bytes":100352,"validate":false},"threads":192}': 'dispatch_kernel_0001', '{"computed_smem_bytes":100352,"constants":[],"ir_name":"flash_kmeans_assign_cleanroom_tcgen05_v10","kwargs":{"smem_bytes":100352,"validate":false},"threads":192}': 'dispatch_kernel_0002', '{"computed_smem_bytes":0,"constants":[],"ir_name":"flash_kmeans_assign_d160_pad192_pack_f9b2_v1","kwargs":{"smem_bytes":0,"validate":false},"threads":256}': 'dispatch_kernel_0004', '{"computed_smem_bytes":149504,"constants":[],"ir_name":"flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1","kwargs":{"smem_bytes":149504,"validate":false},"threads":192}': 'dispatch_kernel_0005', '{"computed_smem_bytes":198656,"constants":[],"ir_name":"flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1","kwargs":{"smem_bytes":198656,"validate":false},"threads":192}': 'dispatch_kernel_0008', '{"computed_smem_bytes":51200,"constants":[],"ir_name":"flash_kmeans_assign_highd_splitd_6fcf_v1","kwargs":{"smem_bytes":51200,"validate":false},"threads":192}': 'dispatch_kernel_0009', '{"computed_smem_bytes":43008,"constants":[],"ir_name":"flash_kmeans_assign_highd_splitk_partial_blockn64_g2r4_b5a6_v1","kwargs":{"smem_bytes":43008,"validate":false},"threads":192}': 'dispatch_kernel_0011', '{"computed_smem_bytes":0,"constants":[],"ir_name":"flash_kmeans_assign_highd_splitk_reduce_blockn64_g2r4_b5a6_v1","kwargs":{"smem_bytes":0,"validate":false},"threads":256}': 'dispatch_kernel_0012', '{"computed_smem_bytes":0,"constants":[],"ir_name":"flash_kmeans_assign_microdim_pack_6cd2_v1","kwargs":{"smem_bytes":0,"validate":false},"threads":256}': 'dispatch_kernel_0013', '{"computed_smem_bytes":51200,"constants":[],"ir_name":"flash_kmeans_assign_microdim_6cd2_v1","kwargs":{"smem_bytes":51200,"validate":false},"threads":192}': 'dispatch_kernel_0014', '{"computed_smem_bytes":0,"constants":[],"ir_name":"flash_kmeans_assign_gap_pad_pack_v1","kwargs":{"smem_bytes":0,"validate":false},"threads":256}': 'dispatch_kernel_0015', '{"computed_smem_bytes":51200,"constants":[],"ir_name":"flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1","kwargs":{"smem_bytes":51200,"validate":false},"threads":192}': 'dispatch_kernel_0016', '{"computed_smem_bytes":133120,"constants":[],"ir_name":"flash_kmeans_assign_cleanroom_tcgen05_v15","kwargs":{"smem_bytes":133120,"validate":false},"threads":192}': 'dispatch_kernel_0048', '{"computed_smem_bytes":43008,"constants":[],"ir_name":"flash_kmeans_assign_highd_splitk_partial_blockn64_g1r4_streamdep_r63_v1","kwargs":{"smem_bytes":43008,"validate":false},"threads":192}': 'dispatch_kernel_0059', '{"computed_smem_bytes":0,"constants":[],"ir_name":"flash_kmeans_assign_highd_splitk_reduce_blockn64_g1r4_streamdep_r63_v1","kwargs":{"smem_bytes":0,"validate":false},"threads":256}': 'dispatch_kernel_0060', '{"computed_smem_bytes":0,"constants":[],"ir_name":"flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1","kwargs":{"smem_bytes":0,"validate":false},"threads":256}': 'dispatch_kernel_0101', '{"computed_smem_bytes":198656,"constants":[],"ir_name":"flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1","kwargs":{"smem_bytes":198656,"validate":false},"threads":192}': 'dispatch_kernel_0109', '{"computed_smem_bytes":75776,"constants":[],"ir_name":"flash_kmeans_assign_highd_paired_xreuse_dualtmem_producer_r47_v1","kwargs":{"smem_bytes":75776,"validate":false},"threads":192}': 'dispatch_kernel_0118', '{"computed_smem_bytes":0,"constants":[],"ir_name":"flash_kmeans_assign_highd_paired_ownerreduce_r39_reduce1_unroll_v1","kwargs":{"smem_bytes":0,"validate":false},"threads":64}': 'dispatch_kernel_0119', '{"computed_smem_bytes":43008,"constants":[],"ir_name":"flash_kmeans_assign_highd_paired_packedpartial_producer_7b3c_v1","kwargs":{"smem_bytes":43008,"validate":false},"threads":192}': 'dispatch_kernel_0121', '{"computed_smem_bytes":0,"constants":[],"ir_name":"flash_kmeans_assign_highd_paired_packedpartial_reduce_r2_7b3c_v1","kwargs":{"smem_bytes":0,"validate":false},"threads":128}': 'dispatch_kernel_0122', '{"computed_smem_bytes":23552,"constants":[],"ir_name":"flash_kmeans_assign_microdim_d16_pipeline4_08f9_v4","kwargs":{"smem_bytes":23552,"validate":false},"threads":192}': 'dispatch_kernel_0126', '{"computed_smem_bytes":52224,"constants":[],"ir_name":"flash_kmeans_assign_microdim_raw_tma_08f9_v1","kwargs":{"smem_bytes":52224,"validate":false},"threads":192}': 'dispatch_kernel_0129', '{"computed_smem_bytes":46592,"constants":[],"ir_name":"flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1","kwargs":{"smem_bytes":46592,"validate":false},"threads":256}': 'dispatch_kernel_0156', '{"computed_smem_bytes":22528,"constants":[],"ir_name":"flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4","kwargs":{"smem_bytes":22528,"validate":false},"threads":192}': 'dispatch_kernel_0229', '{"computed_smem_bytes":26624,"constants":[],"ir_name":"flash_kmeans_assign_d288_exactd_a532_v1","kwargs":{"smem_bytes":26624,"validate":false},"threads":192}': 'dispatch_kernel_0236', '{"computed_smem_bytes":26624,"constants":[],"ir_name":"flash_kmeans_assign_d288_splitk_cta_0438_v1_partial","kwargs":{"smem_bytes":26624,"validate":false},"threads":192}': 'dispatch_kernel_0237', '{"computed_smem_bytes":0,"constants":[],"ir_name":"flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce","kwargs":{"smem_bytes":0,"validate":false},"threads":128}': 'dispatch_kernel_0238', '{"computed_smem_bytes":22528,"constants":[],"ir_name":"flash_kmeans_assign_d480_splitk_partial_d32k256_v1","kwargs":{"smem_bytes":22528,"validate":false},"threads":192}': 'dispatch_kernel_0258', '{"computed_smem_bytes":0,"constants":[],"ir_name":"flash_kmeans_assign_d480_splitk_reduce_d32k256_v1","kwargs":{"smem_bytes":0,"validate":false},"threads":256}': 'dispatch_kernel_0259', '{"computed_smem_bytes":51200,"constants":[],"ir_name":"flash_kmeans_assign_highd_splitk_partial_8de8_v1","kwargs":{"smem_bytes":51200,"validate":false},"threads":192}': 'dispatch_kernel_0277', '{"computed_smem_bytes":0,"constants":[],"ir_name":"flash_kmeans_assign_highd_splitk_reduce_8de8_v1","kwargs":{"smem_bytes":0,"validate":false},"threads":128}': 'dispatch_kernel_0278', '{"computed_smem_bytes":51200,"constants":[],"ir_name":"flash_kmeans_assign_microdim_direct_9c0d_v1","kwargs":{"smem_bytes":51200,"validate":false},"threads":192}': 'dispatch_kernel_0332', '{"computed_smem_bytes":26624,"constants":[],"ir_name":"flash_kmeans_assign_d416_exactd_splitd_a4a579d1_v2","kwargs":{"smem_bytes":26624,"validate":false},"threads":192}': 'dispatch_kernel_0347', '{"computed_smem_bytes":29696,"constants":[],"ir_name":"flash_kmeans_assign_d112_k1024_owner_local_warp_mma_distinct_workfeed_c829_v1","kwargs":{"smem_bytes":29696,"validate":false},"threads":160}': 'dispatch_kernel_0353', '{"computed_smem_bytes":43008,"constants":[],"ir_name":"flash_kmeans_assign_d112_k512_two_owner_peer_only_mma_f826_v1","kwargs":{"smem_bytes":43008,"validate":false},"threads":256}': 'dispatch_kernel_0354', '{"computed_smem_bytes":35840,"constants":[],"ir_name":"flash_kmeans_assign_d112_shared_point_dual_issuer_tcgen05_6e1e_v1","kwargs":{"smem_bytes":35840,"validate":false},"threads":416}': 'dispatch_kernel_0355'} + +import ctypes +import importlib +import json +import threading +import sys +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass, field, replace as _dataclass_replace +from importlib import resources +from types import SimpleNamespace + +from .kernels import get_kernel +from ._runtime import launch_stream_context, resolve_launch_defaults + + +_DISPATCH_LAUNCH_OPTIONS = ContextVar("dispatch_launch_options", default=(None, None)) + + +@contextmanager +def dispatch_launch_options(*, stream=None, timeout_ms=None): + token = _DISPATCH_LAUNCH_OPTIONS.set((stream, timeout_ms)) + try: + yield + finally: + _DISPATCH_LAUNCH_OPTIONS.reset(token) + + +def _resolved_launch_options(stream, timeout_ms): + default_stream, default_timeout_ms = _DISPATCH_LAUNCH_OPTIONS.get() + return ( + default_stream if stream is None else stream, + default_timeout_ms if timeout_ms is None else timeout_ms, + ) + + +_active_launch_capture = ContextVar("flashlib_active_launch_capture", default=None) +_pending_tensor_map_recipe = ContextVar("flashlib_pending_tensor_map_recipe", default=None) +_launch_capture_prepare_lock = threading.RLock() + + +def _replace(value, /, **changes): + replacer = getattr(value, "__replace__", None) + if callable(replacer): + return replacer(**changes) + return _dataclass_replace(value, **changes) + + +dc = SimpleNamespace(replace=_replace) + + +def _import_dispatch_module(short_name): + return importlib.import_module(f"{__package__}._dispatch.{short_name}") + + +_DISPATCH_OWNED_DICT_SUFFIXES = ("CACHE", "SCRATCH", "INPUTS", "OUTPUTS", "FLAGS") + + +def _cache_value_references_owned_object(value, owned_ids, seen): + identity = id(value) + if identity in owned_ids: + return True + if identity in seen: + return False + seen.add(identity) + if isinstance(value, dict): + return any( + _cache_value_references_owned_object(item, owned_ids, seen) + for pair in value.items() + for item in pair + ) + if isinstance(value, (tuple, list, set, frozenset)): + return any( + _cache_value_references_owned_object(item, owned_ids, seen) + for item in value + ) + return False + + +def release_dispatch_caches(owned_objects): + '''Clear route-owned tensor dictionaries after a prepared sequence binds. + + Generated dispatch modules may temporarily cache tensor-map descriptors and + workspaces while a route is prepared. A bound ``PreparedKernelSequence`` + retains every CUDA argument, so those module globals are no longer owners. + To avoid clearing dispatch registries or scalar statistics, this contract + is limited to dict-valued, private, uppercase names with an explicit + workspace/cache suffix below this generated package's ``_dispatch`` + namespace. + ''' + + prefix = f"{__package__}._dispatch." + owned_ids = { + id(value) + for value in owned_objects + if callable(getattr(value, "data_ptr", None)) + } + if not owned_ids: + return 0 + cleared = 0 + for module_name, module in tuple(sys.modules.items()): + if module is None or not module_name.startswith(prefix): + continue + for name, value in tuple(vars(module).items()): + if ( + name.startswith("_") + and name.endswith(_DISPATCH_OWNED_DICT_SUFFIXES) + and name.isupper() + and isinstance(value, dict) + ): + removed = False + for key, item in tuple(value.items()): + if _cache_value_references_owned_object(item, owned_ids, set()): + value.pop(key, None) + removed = True + cleared += int(removed) + return cleared + + +def _decode_capture(value): + if isinstance(value, dict) and "__ir__" in value: + return _ir_proxy( + value["__ir__"], + value.get("threads", 256), + value.get("computed_smem_bytes", 0), + value.get("cluster_dims", (1, 1, 1)), + value.get("cta_group", 1), + value.get("constants", ()), + value.get("arg_keys", ()), + ) + if isinstance(value, dict) and set(value) == {"__kernel__"}: + return DispatchKernel(value["__kernel__"]) + if isinstance(value, dict) and set(value) == {"__kernel_source__"}: + return value["__kernel_source__"] + if isinstance(value, dict) and set(value) == {"__tuple__"}: + return tuple(_decode_capture(item) for item in value["__tuple__"]) + if isinstance(value, dict) and set(value) == {"__dict_items__"}: + return { + _decode_capture(key): _decode_capture(item) + for key, item in value["__dict_items__"] + } + if isinstance(value, dict): + return {key: _decode_capture(item) for key, item in value.items()} + if isinstance(value, list): + return [_decode_capture(item) for item in value] + return value + + +@dataclass(frozen=True) +class _IRProxy: + symbol: str + threads: int = 256 + computed_smem_bytes: int = 0 + constants: tuple = () + grid: object = None + arg_keys: tuple = () + + def __replace__(self, /, **changes): + values = { + "symbol": self.symbol, + "threads": self.threads, + "computed_smem_bytes": self.computed_smem_bytes, + "constants": self.constants, + "grid": self.grid, + "arg_keys": self.arg_keys, + } + unknown = sorted(set(changes) - set(values)) + if unknown: + raise TypeError(f"unknown frozen WeaveIR field(s): {unknown}") + values.update(changes) + return _IRProxy(**values) + + +def _ir_proxy( + name, threads=256, computed_smem_bytes=0, cluster_dims=(1, 1, 1), + cta_group=1, constants=(), arg_keys=(), +): + return _IRProxy( + name.rpartition(":")[2], int(threads), int(computed_smem_bytes), + tuple(tuple(item) for item in constants), + SimpleNamespace(cluster_dims=tuple(cluster_dims), cta_group=int(cta_group)), + tuple(arg_keys), + ) + + +def pack_kernel_args(schedule, /, **bindings): + expected = tuple(schedule.arg_keys) + missing = sorted(set(expected) - set(bindings)) + unexpected = sorted(set(bindings) - set(expected)) + if missing or unexpected: + raise ValueError( + f"kernel argument bindings do not match frozen WeaveIR.args: " + f"missing={missing!r}, unexpected={unexpected!r}" + ) + return [bindings[key] for key in expected] + + +class PreparedKernelSequence: + def __init__( + self, + launches, + result, + input_bindings=(), + result_template=None, + tensor_map_bindings=(), + input_alias_topology=(), + stream=None, + ): + if not launches: + raise RuntimeError("prepared semantic route did not capture a CUDA launch") + self._launches = tuple(launches) + self._result = result + self._input_bindings = tuple(tuple(bindings) for bindings in input_bindings) + if self._input_bindings and len(self._input_bindings) != len(self._launches): + raise RuntimeError("prepared semantic route has corrupt input bindings") + self._result_template = result_template + self._input_alias_topology = tuple( + tuple(group) for group in input_alias_topology + ) + direct_input_keys = {key for bindings in self._input_bindings for _, key in bindings} + self._direct_input_keys = tuple(sorted(direct_input_keys)) + self._tensor_map_bindings = _own_tensor_map_bindings( + self._launches, + tuple(tensor_map_bindings), + stream=stream, + ) + self._bound_input_keys = tuple( + sorted(direct_input_keys | {binding.input_key for binding in self._tensor_map_bindings}) + ) + self._input_references_retained = True + + @property + def launch_count(self): + return len(self._launches) + + @property + def bound_input_keys(self): + return self._bound_input_keys + + def rebind_inputs( + self, + inputs, + *, + stream=None, + materialize_result=True, + preserve_prepared_stream=False, + retain_input_references=True, + ): + if not isinstance(materialize_result, bool): + raise TypeError("materialize_result must be a bool") + if not isinstance(preserve_prepared_stream, bool): + raise TypeError("preserve_prepared_stream must be a bool") + if not isinstance(retain_input_references, bool): + raise TypeError("retain_input_references must be a bool") + if preserve_prepared_stream and stream is not None: + raise ValueError("preserve_prepared_stream requires stream=None") + if materialize_result and not retain_input_references: + raise ValueError( + "materialize_result requires retain_input_references so the result remains valid" + ) + if not any(self._input_bindings) and not self._tensor_map_bindings: + raise RuntimeError( + "prepared semantic route has no input bindings; " + "capture it with capture_kernel_launches(inputs=...)" + ) + missing = sorted(set(self.bound_input_keys) - set(inputs)) + if missing: + raise KeyError(f"missing prepared semantic input binding(s): {missing!r}") + _validate_public_tensor_alias_topology(inputs, self._input_alias_topology) + pointer_values = None + inputs_already_scrubbed = False + if not retain_input_references: + pointer_values = {} + for key in self._direct_input_keys: + value = inputs[key] + data_ptr = getattr(value, "data_ptr", None) + if not callable(data_ptr): + raise TypeError(f"prepared CUDA tensor binding {key!r} is not tensor-like") + pointer_values[key] = int(data_ptr()) + inputs_already_scrubbed = not self._input_references_retained + with launch_stream_context(stream): + for binding in self._tensor_map_bindings: + binding.refresh(inputs[binding.input_key]) + for launch, bindings in zip(self._launches, self._input_bindings, strict=True): + launch.rebind_tensor_arguments( + bindings, + inputs, + stream=stream, + preserve_stream=preserve_prepared_stream, + retain_inputs=retain_input_references, + pointer_values=pointer_values, + inputs_already_scrubbed=inputs_already_scrubbed, + ) + self._input_references_retained = retain_input_references + # Stateful public runtimes may own the output independently of the + # semantic return tree. Let those callers skip recursively rebuilding + # a result they will not observe while retaining the default behavior + # for normal prepared-dispatch callers. + if materialize_result and self._result_template is not None: + self._result = _materialize_result_template(self._result_template, inputs) + return self + + def _rebind_stream_bound_scrubbed_inputs(self, inputs, *, stream): + '''Rebind one validated fixed-stream runtime slot without generic checks. + + This private path is valid only after a stateful wrapper selected the + sequence through a cache key containing the complete public pointer + alias topology, recorded every caller-owned tensor, and scrubbed the + sequence's caller references. Public prepared callers continue to use + :meth:`rebind_inputs` and its full validation. + ''' + if self._input_references_retained: + raise RuntimeError( + "fixed-stream semantic rebind requires scrubbed input references" + ) + if stream is None: + raise ValueError("fixed-stream semantic rebind requires an explicit stream") + pointer_values = {} + for key in self._direct_input_keys: + value = inputs[key] + data_ptr = getattr(value, "data_ptr", None) + if not callable(data_ptr): + raise TypeError(f"prepared CUDA tensor binding {key!r} is not tensor-like") + pointer_values[key] = int(data_ptr()) + for binding in self._tensor_map_bindings: + binding.rebind_stream_bound(inputs[binding.input_key], stream=stream) + for launch, bindings in zip(self._launches, self._input_bindings, strict=True): + launch.rebind_tensor_arguments( + bindings, + inputs, + preserve_stream=True, + retain_inputs=False, + pointer_values=pointer_values, + inputs_already_scrubbed=True, + ) + return self + + def release_bound_inputs(self): + '''Drop caller tensor references after their launch stream was recorded.''' + if self._input_references_retained: + for launch, bindings in zip(self._launches, self._input_bindings, strict=True): + keepalive = list(launch._keepalive) + for index, _key in bindings: + value = keepalive[index] + data_ptr = getattr(value, "data_ptr", None) + if callable(data_ptr): + keepalive[index] = int(data_ptr()) + launch._keepalive = tuple(keepalive) + self._input_references_retained = False + self._result = None + + def record_stream(self, stream): + '''Tie every tensor launch argument, including private scratch, to a stream.''' + if stream is None: + raise ValueError("prepared semantic record_stream requires an explicit stream") + seen = set() + for launch in self._launches: + for value in launch._keepalive: + identity = id(value) + record_stream = getattr(value, "record_stream", None) + if identity not in seen and callable(record_stream): + seen.add(identity) + record_stream(stream) + # Variant-bank descriptor tensors are slot-owned but only the active + # variant appears in a launch keepalive; record every variant so a + # non-synchronizing release stays allocator-safe. + for binding in self._tensor_map_bindings: + for value in binding.variants.values(): + identity = id(value) + record_stream = getattr(value, "record_stream", None) + if identity not in seen and callable(record_stream): + seen.add(identity) + record_stream(stream) + + def _finish_rebind(self, result): + self._result = result + return self + + def __call__(self, _inputs=None, *, stream=None, timeout_ms=None): + last = len(self._launches) - 1 + for index, launch in enumerate(self._launches): + launch.launch(stream=stream, timeout_ms=timeout_ms if index == last else None) + return self._result + + +class KernelLaunchCapture: + def __init__(self, *, stream=None, arch=None, inputs=None, rebind=None): + if rebind is not None and not isinstance(rebind, PreparedKernelSequence): + raise TypeError("rebind must be a PreparedKernelSequence") + if inputs is not None and rebind is not None: + raise ValueError("inputs and rebind are mutually exclusive capture modes") + if rebind is not None: + raise RuntimeError( + "capture(rebind=...) is unsupported because an in-place topology " + "update cannot be transactional; capture a new sequence instead" + ) + self.stream = stream + self.arch = arch + self._launches = [] + self._input_bindings = [] + self._input_key_by_identity = _public_tensor_input_identities(inputs) + self._input_key_by_pointer = _public_tensor_input_pointers(inputs) + self._input_alias_topology = _public_tensor_alias_topology(inputs) + self._tensor_map_bindings = {} + self._route_caches_released = False + self._rebind = rebind + self._rebind_index = 0 + self.host_data_reads = 0 + + @property + def host_data_dependent(self): + '''True when the route read device memory while its kernels were only + being recorded — its host branch decisions cannot be frozen.''' + return self.host_data_reads > 0 + + @property + def rebinding(self): + return self._rebind is not None + + def add(self, launch): + if self.rebinding: + raise RuntimeError("rebind capture requires launch topology, not a newly prepared launch") + self._launches.append(launch) + self._input_bindings.append( + () + if not self._input_key_by_identity + else tuple( + (index, self._input_key_by_identity[id(arg)]) + for index, arg in enumerate(launch._keepalive) + if id(arg) in self._input_key_by_identity + ) + ) + for arg in launch._keepalive: + recipe = getattr(arg, "_loom_tensor_map_recipe", None) + if recipe is None: + continue + source_pointer = int(recipe[2]) + input_key = self._input_key_by_pointer.get(source_pointer) + if input_key is None: + continue + self._tensor_map_bindings.setdefault( + id(arg), + _TensorMapBinding( + input_key=input_key, + tensor=arg, + recipe=tuple(recipe), + pointer=source_pointer, + ), + ) + + def add_kernel_launch( + self, + exported, + *, + mode, + grid, + block, + args, + arg_types, + shared_mem, + stream, + cluster_dims=None, + ): + resolved_arch, resolved_stream, _ = resolve_launch_defaults( + arch=self.arch, + stream=self.stream if self.stream is not None else stream, + timeout_ms=None, + ) + with launch_stream_context(resolved_stream): + kernel = exported.compile(arch=resolved_arch, options=["--use_fast_math"]) + kwargs = { + "grid": grid, + "block": block, + "args": tuple(args), + "arg_types": arg_types, + "shared_mem": shared_mem, + "stream": resolved_stream, + } + if self.rebinding: + if self._rebind_index >= self._rebind.launch_count: + raise RuntimeError( + "prepared semantic route launch-count mismatch: " + f"expected {self._rebind.launch_count}, captured more launches" + ) + prepared = self._rebind._launches[self._rebind_index] + if mode == "cluster": + kernel.rebind_launch_cluster( + prepared, cluster_dims=cluster_dims, **kwargs + ) + elif mode == "cooperative": + kernel.rebind_launch_cooperative(prepared, **kwargs) + elif mode == "regular": + kernel.rebind_launch(prepared, **kwargs) + else: + raise RuntimeError(f"unsupported captured launch mode: {mode!r}") + self._rebind_index += 1 + return + if mode == "cluster": + prepared = kernel.prepare_launch_cluster( + cluster_dims=cluster_dims, **kwargs + ) + elif mode == "cooperative": + prepared = kernel.prepare_launch_cooperative(**kwargs) + elif mode == "regular": + prepared = kernel.prepare_launch(**kwargs) + else: + raise RuntimeError(f"unsupported captured launch mode: {mode!r}") + self.add(prepared) + + def bind(self, result): + if self.rebinding: + if self._rebind_index != self._rebind.launch_count: + raise RuntimeError( + "prepared semantic route launch-count mismatch: " + f"expected {self._rebind.launch_count}, captured {self._rebind_index}" + ) + return self._rebind._finish_rebind(result) + result_template = _capture_result_template(result, self._input_key_by_identity) + route_cache_owned_objects = self._route_cache_owned_objects() + sequence = PreparedKernelSequence( + self._launches, + result, + self._input_bindings, + result_template, + tuple(self._tensor_map_bindings.values()), + self._input_alias_topology, + self.stream, + ) + self.release_route_caches(route_cache_owned_objects) + return sequence + + def _route_cache_owned_objects(self): + return tuple(arg for launch in self._launches for arg in launch._keepalive) + + def release_route_caches(self, owned_objects=None): + if self._route_caches_released: + return 0 + if owned_objects is None: + owned_objects = self._route_cache_owned_objects() + self._route_caches_released = True + return release_dispatch_caches(tuple(owned_objects)) + + +@dataclass(frozen=True) +class _BoundInputResult: + key: str + + +@dataclass +class _TensorMapBinding: + input_key: str + tensor: object + recipe: tuple + pointer: int + pointer_carriers: tuple = () + staging_slots: list = field(default_factory=list) + variants: dict = field(default_factory=dict) + variant_capacity: int = 4 + + def __post_init__(self): + if not self.variants: + self.variants[self.pointer] = self.tensor + + def _acquire_staging_slot(self, torch): + for slot in self.staging_slots: + if slot.event.query(): + return slot + slot = _TensorMapStagingSlot( + host_buffer=torch.empty(128, dtype=torch.uint8, pin_memory=True), + event=torch.cuda.Event(blocking=False, interprocess=False), + ) + self.staging_slots.append(slot) + return slot + + def _encode_into(self, pointer, tensor, *, stream=None): + from cuda.bindings import driver + import torch + + arguments = list(self.recipe) + arguments[2] = pointer + err, tmap = driver.cuTensorMapEncodeTiled(*arguments) + if err != 0: + raise RuntimeError(f"cuTensorMapEncodeTiled rebind failed: CUresult={err}") + slot = self._acquire_staging_slot(torch) + ctypes.memmove( + int(slot.host_buffer.data_ptr()), + int(tmap.getPtr()), + 128, + ) + tensor.copy_(slot.host_buffer, non_blocking=True) + slot.event.record(torch.cuda.current_stream() if stream is None else stream) + self.tensor = tensor + self.pointer = pointer + self.recipe = tuple(arguments) + tensor._loom_tensor_map_recipe = self.recipe + + def _activate(self, pointer, tensor): + descriptor_pointer = int(tensor.data_ptr()) + for carrier in self.pointer_carriers: + carrier.value = descriptor_pointer + arguments = list(self.recipe) + arguments[2] = pointer + self.tensor = tensor + self.pointer = pointer + self.recipe = tuple(arguments) + tensor._loom_tensor_map_recipe = self.recipe + + def refresh(self, source, *, stream=None): + pointer = int(source.data_ptr()) + if pointer == self.pointer: + return + self._encode_into(pointer, self.tensor, stream=stream) + # Generic public rebinding mutates the active descriptor in place. + # Reset the private variant bank so no stale pointer key can name it. + self.variants.clear() + self.variants[pointer] = self.tensor + + def rebind_stream_bound(self, source, *, stream): + if stream is None: + raise ValueError("stream-bound tensor-map rebind requires an explicit stream") + pointer = int(source.data_ptr()) + if pointer == self.pointer: + return + cached = self.variants.pop(pointer, None) + if cached is not None: + # LRU recency: re-insert the hit so eviction removes the + # least-recently-activated variant, not the newest one. + self.variants[pointer] = cached + self._activate(pointer, cached) + return + + import torch + + if len(self.variants) < self.variant_capacity: + tensor = torch.empty_like(self.tensor) + record_stream = getattr(tensor, "record_stream", None) + if callable(record_stream): + record_stream(stream) + else: + tensor = self.variants.pop(next(iter(self.variants))) + self._encode_into(pointer, tensor, stream=stream) + self.variants[pointer] = tensor + self._activate(pointer, tensor) + + +@dataclass +class _TensorMapStagingSlot: + host_buffer: object + event: object + + +def _own_tensor_map_bindings(launches, bindings, *, stream): + '''Clone cached descriptors and patch every launch to slot-owned storage.''' + + if not bindings: + return () + owned_by_identity = {} + owned_bindings = [] + with launch_stream_context(stream): + for binding in bindings: + original = binding.tensor + owned = original.clone() + owned._loom_tensor_map_recipe = binding.recipe + metadata = getattr(original, "_loom_tma_metadata", None) + if metadata is not None: + owned._loom_tma_metadata = metadata + owned_by_identity[id(original)] = owned + pointer_carriers = tuple( + launch._packed._prevent_gc[index] + for launch in launches + for index, arg in enumerate(launch._keepalive) + if id(arg) == id(original) + ) + if not pointer_carriers or any( + type(carrier) is not ctypes.c_void_p for carrier in pointer_carriers + ): + raise RuntimeError("captured tensor-map binding has invalid pointer carriers") + owned_bindings.append( + _TensorMapBinding( + input_key=binding.input_key, + tensor=owned, + recipe=binding.recipe, + pointer=binding.pointer, + pointer_carriers=pointer_carriers, + ) + ) + for launch in launches: + replacements = { + index: owned_by_identity[id(arg)] + for index, arg in enumerate(launch._keepalive) + if id(arg) in owned_by_identity + } + if replacements: + launch.rebind_arguments(replacements, stream=stream) + return tuple(owned_bindings) + + +def _public_tensor_input_identities(inputs): + if inputs is None: + return {} + if not hasattr(inputs, "items"): + raise TypeError("capture inputs must be a mapping") + identities = {} + for key, value in inputs.items(): + if ( + isinstance(key, str) + and not key.startswith("_") + and callable(getattr(value, "data_ptr", None)) + ): + identities.setdefault(id(value), key) + return identities + + +def _public_tensor_input_pointers(inputs): + if inputs is None: + return {} + pointers = {} + for key, value in inputs.items(): + if ( + isinstance(key, str) + and not key.startswith("_") + and callable(getattr(value, "data_ptr", None)) + ): + pointers.setdefault(int(value.data_ptr()), key) + return pointers + + +def _public_tensor_alias_topology(inputs, keys=None): + '''Return the complete pointer-equality partition of public tensor inputs.''' + + if inputs is None: + return () + if not hasattr(inputs, "items"): + raise TypeError("capture inputs must be a mapping") + if keys is None: + selected = [ + key + for key, value in inputs.items() + if ( + isinstance(key, str) + and not key.startswith("_") + and callable(getattr(value, "data_ptr", None)) + ) + ] + else: + selected = list(keys) + missing = sorted(set(selected) - set(inputs)) + if missing: + raise KeyError(f"missing prepared semantic alias binding(s): {missing!r}") + invalid = sorted( + key + for key in selected + if not callable(getattr(inputs[key], "data_ptr", None)) + ) + if invalid: + raise TypeError( + f"prepared semantic alias binding(s) must be tensor-like: {invalid!r}" + ) + groups = {} + for key in selected: + groups.setdefault(int(inputs[key].data_ptr()), []).append(key) + return tuple(sorted(tuple(sorted(group)) for group in groups.values())) + + +def _validate_public_tensor_alias_topology(inputs, expected): + if not expected: + return + keys = tuple(key for group in expected for key in group) + actual = _public_tensor_alias_topology(inputs, keys) + if actual != expected: + raise RuntimeError( + "prepared semantic public tensor alias topology changed: " + f"expected {expected!r}, got {actual!r}; capture a new sequence" + ) + + +def _capture_result_template(value, input_key_by_identity): + key = input_key_by_identity.get(id(value)) + if key is not None: + return _BoundInputResult(key) + if isinstance(value, tuple): + return tuple(_capture_result_template(item, input_key_by_identity) for item in value) + if isinstance(value, list): + return [_capture_result_template(item, input_key_by_identity) for item in value] + if isinstance(value, dict): + return { + key: _capture_result_template(item, input_key_by_identity) + for key, item in value.items() + } + return value + + +def _materialize_result_template(value, inputs): + if isinstance(value, _BoundInputResult): + return inputs[value.key] + if isinstance(value, tuple): + return tuple(_materialize_result_template(item, inputs) for item in value) + if isinstance(value, list): + return [_materialize_result_template(item, inputs) for item in value] + if isinstance(value, dict): + return {key: _materialize_result_template(item, inputs) for key, item in value.items()} + return value + + +@contextmanager +def capture_kernel_launches(*, stream=None, arch=None, inputs=None, rebind=None): + import torch + + with _launch_capture_prepare_lock: + if _active_launch_capture.get() is not None: + raise RuntimeError("nested kernel launch capture is not supported") + capture = KernelLaunchCapture(stream=stream, arch=arch, inputs=inputs, rebind=rebind) + token = _active_launch_capture.set(capture) + # Captured launches are marshalled, not run, so any device read the + # route performs mid-traversal (for example an ``overflow_flag.item()`` + # certification) observes memory its recorded kernels never wrote. + # Interpose ``torch.Tensor.item`` for the capture's duration and count + # CUDA-tensor reads; a nonzero count marks the capture + # ``host_data_dependent`` so the plan builder keeps that signature on + # the per-call launcher instead of freezing an unreproducible branch. + # Captures are serialized by ``_launch_capture_prepare_lock``, so the + # process-global interpose cannot nest. ``.item()`` is the only device + # read the vendored dispatch modules perform on their launch paths. + # Torch test doubles without a ``Tensor.item`` skip the interpose. + original_tensor_item = getattr(getattr(torch, "Tensor", None), "item", None) + + def _observed_item(tensor): + if getattr(tensor, "is_cuda", False): + capture.host_data_reads += 1 + return original_tensor_item(tensor) + + if original_tensor_item is not None: + torch.Tensor.item = _observed_item + try: + if stream is None: + yield capture + else: + with torch.cuda.stream(stream): + yield capture + finally: + if original_tensor_item is not None: + torch.Tensor.item = original_tensor_item + capture.release_route_caches() + _active_launch_capture.reset(token) + + +class DispatchKernel: + def __init__(self, alias, symbol=None): + self.exported = get_kernel(alias) + self.symbol = symbol or self.exported.spec.symbol + + def __enter__(self): + return self + + def __exit__(self, *args): + return None + + def launch(self, *, grid, block, args, shared_mem=0, stream=None, timeout_ms=None, **kwargs): + stream, timeout_ms = _resolved_launch_options(stream, timeout_ms) + capture = _active_launch_capture.get() + if capture is not None: + capture.add_kernel_launch( + self.exported, + mode="regular", + grid=grid, + block=block, + args=args, + arg_types=self.exported.arg_types, + shared_mem=shared_mem, + stream=stream, + ) + return + self.exported.launch( + *args, grid=grid, block=block, shared_mem=shared_mem, stream=stream, + timeout_ms=timeout_ms, options=["--use_fast_math"], + ) + + def launch_cluster( + self, *, grid, block, args, cluster_dims, shared_mem=0, stream=None, + timeout_ms=None, **kwargs + ): + stream, timeout_ms = _resolved_launch_options(stream, timeout_ms) + capture = _active_launch_capture.get() + if capture is not None: + capture.add_kernel_launch( + self.exported, + mode="cluster", + grid=grid, + block=block, + args=args, + arg_types=self.exported.arg_types, + cluster_dims=cluster_dims, + shared_mem=shared_mem, + stream=stream, + ) + return + arch, stream, timeout_ms = resolve_launch_defaults( + arch=None, + stream=stream, + timeout_ms=timeout_ms, + ) + with launch_stream_context(stream): + kernel = self.exported.compile(arch=arch, options=["--use_fast_math"]) + kernel.launch_cluster( + grid=grid, block=block, args=tuple(args), + arg_types=self.exported.arg_types, + cluster_dims=cluster_dims, shared_mem=shared_mem, stream=stream, + timeout_ms=timeout_ms, + ) + + def launch_cooperative( + self, *, grid, block, args, shared_mem=0, stream=None, timeout_ms=None, **kwargs + ): + stream, timeout_ms = _resolved_launch_options(stream, timeout_ms) + capture = _active_launch_capture.get() + if capture is not None: + capture.add_kernel_launch( + self.exported, + mode="cooperative", + grid=grid, + block=block, + args=args, + arg_types=self.exported.arg_types, + shared_mem=shared_mem, + stream=stream, + ) + return + arch, stream, timeout_ms = resolve_launch_defaults( + arch=None, + stream=stream, + timeout_ms=timeout_ms, + ) + with launch_stream_context(stream): + kernel = self.exported.compile(arch=arch, options=["--use_fast_math"]) + kernel.launch_cooperative( + grid=grid, block=block, args=tuple(args), + arg_types=self.exported.arg_types, + shared_mem=shared_mem, stream=stream, timeout_ms=timeout_ms, + ) + + def prepare_launch(self, **kwargs): + return _PreparedDispatchLaunch(self.launch, kwargs) + + def prepare_launch_cluster(self, **kwargs): + return _PreparedDispatchLaunch(self.launch_cluster, kwargs) + + def prepare_launch_cooperative(self, **kwargs): + return _PreparedDispatchLaunch(self.launch_cooperative, kwargs) + + +class _PreparedDispatchLaunch: + def __init__(self, launch, kwargs): + self._launch = launch + self._kwargs = dict(kwargs) + + def launch(self, timeout_ms=None): + kwargs = dict(self._kwargs) + if timeout_ms is not None: + kwargs["timeout_ms"] = timeout_ms + return self._launch(**kwargs) + + +CUDAKernel = DispatchKernel + + +def compile_cuda(source, **kwargs): + return source + + +def detect_gpu_arch(): + import torch + major, minor = torch.cuda.get_device_capability() + return f"sm_{major}{minor}a" + + +def _cuda_include_dirs(): + return [] + + +def arch_flag_for_cc(major, minor): + sm = int(major) * 10 + int(minor) + return f"sm_{sm}a" if sm >= 90 else f"sm_{sm}" + + +def _capture_cuTensorMapEncodeTiled(*arguments): + '''Encode a tensor map while retaining a pointer-rebind recipe.''' + from cuda.bindings import driver + + result = driver.cuTensorMapEncodeTiled(*arguments) + if result[0] == 0: + _pending_tensor_map_recipe.set(tuple(arguments)) + else: + _pending_tensor_map_recipe.set(None) + return result + + +def _tmap_to_device(tmap, metadata=None): + import torch + del metadata + recipe = _pending_tensor_map_recipe.get() + _pending_tensor_map_recipe.set(None) + host_ptr = tmap.getPtr() + raw = bytes((ctypes.c_ubyte * 128).from_address(host_ptr)) + host = torch.frombuffer(bytearray(raw), dtype=torch.uint8) + device = torch.empty(128, dtype=torch.uint8, device="cuda") + device.copy_(host) + if recipe is not None: + device._loom_tensor_map_recipe = recipe + return device + + +class Swizzle: + # Standalone spellings used only by stripped TMA metadata helpers. + SZ_128B = "128B" + SZ_64B = "64B" + SZ_32B = "32B" + NONE = "none" + + +class TensorMapMetadata: + # Compatibility carrier; frozen launch packing needs no metadata. + def __init__(self, **values): + self.__dict__.update(values) + + +def attach_tma_metadata(tensor, metadata): + tensor._loom_tensor_map_metadata = metadata + return tensor + + +def create_tensor_map(data_ptr, dim0, dim1, box0, box1, stride1_bytes): + '''Create a rank-2 BF16 128B-swizzled map with a rebind recipe.''' + from cuda.bindings import driver + + err, tmap = _capture_cuTensorMapEncodeTiled( + driver.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, + 2, + data_ptr, + [driver.cuuint64_t(dim0), driver.cuuint64_t(dim1)], + [driver.cuuint64_t(stride1_bytes)], + [driver.cuuint32_t(box0), driver.cuuint32_t(box1)], + [driver.cuuint32_t(1), driver.cuuint32_t(1)], + driver.CUtensorMapInterleave.CU_TENSOR_MAP_INTERLEAVE_NONE, + driver.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B, + driver.CUtensorMapL2promotion.CU_TENSOR_MAP_L2_PROMOTION_NONE, + driver.CUtensorMapFloatOOBfill.CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE, + ) + if err != 0: + raise RuntimeError(f"cuTensorMapEncodeTiled (2D BF16) failed: CUresult={err}") + return _tmap_to_device(tmap) + + +def _create_tensor_map_3d(data_ptr, global_height, shared_height, width, block_width, swizzle): + from cuda.bindings import driver + atoms = {"128B": 64, "64B": 32, "32B": 16} + swizzles = { + "128B": driver.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B, + "64B": driver.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_64B, + "32B": driver.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_32B, + } + try: + atom = atoms[swizzle] + swizzle_value = swizzles[swizzle] + except KeyError as exc: + raise ValueError(f"unsupported 3D tensor-map swizzle: {swizzle}") from exc + err, tmap = _capture_cuTensorMapEncodeTiled( + driver.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, + 3, + data_ptr, + [driver.cuuint64_t(atom), driver.cuuint64_t(global_height), driver.cuuint64_t(width // atom)], + [driver.cuuint64_t(width * 2), driver.cuuint64_t(atom * 2)], + [driver.cuuint32_t(atom), driver.cuuint32_t(shared_height), driver.cuuint32_t(block_width // atom)], + [driver.cuuint32_t(1), driver.cuuint32_t(1), driver.cuuint32_t(1)], + driver.CUtensorMapInterleave.CU_TENSOR_MAP_INTERLEAVE_NONE, + swizzle_value, + driver.CUtensorMapL2promotion.CU_TENSOR_MAP_L2_PROMOTION_NONE, + driver.CUtensorMapFloatOOBfill.CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE, + ) + if err != 0: + raise RuntimeError(f"cuTensorMapEncodeTiled failed: CUresult={err}") + return _tmap_to_device(tmap) + + +def create_tensor_map_3d(data_ptr, global_height, shared_height, width, block_width): + return _create_tensor_map_3d(data_ptr, global_height, shared_height, width, block_width, "128B") + + +def create_tensor_map_3d_64b(data_ptr, global_height, shared_height, width, block_width): + return _create_tensor_map_3d(data_ptr, global_height, shared_height, width, block_width, "64B") + + +def create_tensor_map_3d_32b(data_ptr, global_height, shared_height, width, block_width): + return _create_tensor_map_3d(data_ptr, global_height, shared_height, width, block_width, "32B") + + +def generate_kernel(ir, **kwargs): + request_key = json.dumps( + { + "ir_name": ir.symbol, + "constants": [[str(name), value] for name, value in ir.constants], + "threads": int(ir.threads), + "computed_smem_bytes": int(ir.computed_smem_bytes), + "kwargs": kwargs, + }, + sort_keys=True, separators=(",", ":"), default=repr, + ) + alias = _KERNEL_ALIAS_BY_REQUEST.get(request_key) + if alias is None: + alias = _KERNEL_ALIAS_BY_IR_NAME.get(ir.symbol) + if alias is None: + raise RuntimeError(f"uncaptured dispatcher specialization for {ir.symbol}") + return alias + + +def generate_kernel_bundle(*args, **kwargs): + raise RuntimeError("uncaptured dispatcher bundle specialization") + + +def _all_shapes(): + package = __package__ or __name__.rpartition(".")[0] + text = resources.files(package).joinpath("_dispatch_shapes.json").read_text(encoding="utf-8") + return json.loads(text) + + +class _CanonicalShapes: + '''Lazy contract-shape view backed by the exported plan ledger.''' + + def __iter__(self): + return iter(_all_shapes()) + + def __len__(self): + return len(_all_shapes()) + + def __getitem__(self, index): + return _all_shapes()[index] + + +CANONICAL_SHAPES = _CanonicalShapes() + + +def select_named_shapes(labels): + labels = [labels] if isinstance(labels, str) else list(labels) + by_label = {row["label"]: row for row in _all_shapes()} + return [by_label[label] for label in labels] + + +def evaluate(*args, **kwargs): + raise RuntimeError("Cake eval harness is not part of the standalone runtime") diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch_shapes.json b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch_shapes.json new file mode 100644 index 00000000..cd982b94 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_dispatch_shapes.json @@ -0,0 +1,4458 @@ +[ + { + "label": "d895_expanded_heldout_neighborhood_d80_b2_n2176_k1024_d80", + "params": { + "B": 2, + "D": 80, + "K": 1024, + "N": 2176, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.7696, + "evolution_flashlib_ms": 0.402944, + "evolution_kernel_ms": 0.176928, + "evolution_speedup": 2.2774, + "evolution_tflops": 4.0301, + "route": "lowdim_e50c_v1", + "source": "expanded9" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d895_expanded_guard_miss_fallback_d128_b1_n1408_k512_d128", + "params": { + "B": 1, + "D": 128, + "K": 512, + "N": 1408, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.53, + "evolution_flashlib_ms": 0.348175, + "evolution_kernel_ms": 0.144544, + "evolution_speedup": 2.4088, + "evolution_tflops": 1.2768, + "route": "aligned_v10_fallback", + "source": "expanded9" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d895_expanded_forced_fallback_d128_b1_n1664_k256_d128", + "params": { + "B": 1, + "D": 128, + "K": 256, + "N": 1664, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.3924, + "evolution_flashlib_ms": 0.277904, + "evolution_kernel_ms": 0.142032, + "evolution_speedup": 1.9566, + "evolution_tflops": 0.7678, + "route": "aligned_v10_fallback", + "source": "expanded9" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d895_expanded_guard_boundary_d144_b1_n128_k256_d144", + "params": { + "B": 1, + "D": 144, + "K": 256, + "N": 128, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0205, + "evolution_flashlib_ms": 0.459647, + "evolution_kernel_ms": 0.1664, + "evolution_speedup": 2.7623, + "evolution_tflops": 0.0567, + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "source": "expanded9" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d895_expanded_tail_divisibility_d176_b1_n1152_k512_d176", + "params": { + "B": 1, + "D": 176, + "K": 512, + "N": 1152, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.5253, + "evolution_flashlib_ms": 0.395232, + "evolution_kernel_ms": 0.174944, + "evolution_speedup": 2.2592, + "evolution_tflops": 1.1868, + "route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "source": "expanded9" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d895_expanded_guard_overlap_d192_b1_n1024_k768_d192", + "params": { + "B": 1, + "D": 192, + "K": 768, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.6618, + "evolution_flashlib_ms": 0.45632, + "evolution_kernel_ms": 0.17168, + "evolution_speedup": 2.658, + "evolution_tflops": 1.759, + "route": "d192_paired_repeated_mma_v1", + "source": "expanded9" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d895_expanded_random_legal_d320_b1_n1280_k512_d320", + "params": { + "B": 1, + "D": 320, + "K": 512, + "N": 1280, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.826, + "evolution_flashlib_ms": 0.507759, + "evolution_kernel_ms": 0.152496, + "evolution_speedup": 3.3297, + "evolution_tflops": 2.7504, + "route": "highd_splitd_single_tile_6fcf_v1", + "source": "expanded9" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d895_expanded_request_specific_d448_b1_n1280_k1024_d448", + "params": { + "B": 1, + "D": 448, + "K": 1024, + "N": 1280, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.7411, + "evolution_flashlib_ms": 0.674512, + "evolution_kernel_ms": 0.173743, + "evolution_speedup": 3.8822, + "evolution_tflops": 6.7594, + "route": "highd_splitd_single_tile_6fcf_v1", + "source": "expanded9" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d895_expanded_request_specific_d512_splitk_b1_n512_k8192_d512", + "params": { + "B": 1, + "D": 512, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1375, + "evolution_flashlib_ms": 31.230383, + "evolution_kernel_ms": 0.257392, + "evolution_speedup": 121.3339, + "evolution_tflops": 16.6865, + "route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "source": "expanded9" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d16_paired_b2_n4096_k1024_d16", + "params": { + "B": 2, + "D": 16, + "K": 1024, + "N": 4096, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.8529, + "evolution_flashlib_ms": 0.314751, + "evolution_kernel_ms": 0.175568, + "evolution_speedup": 1.7928, + "evolution_tflops": 1.529, + "route": "microdim_hybrid_9c0d_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d16_fallback_b3_n2432_k512_d16", + "params": { + "B": 3, + "D": 16, + "K": 512, + "N": 2432, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.4098, + "evolution_flashlib_ms": 0.291712, + "evolution_kernel_ms": 0.127456, + "evolution_speedup": 2.2887, + "evolution_tflops": 0.9379, + "route": "microdim_hybrid_9c0d_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d16_small_b4_n1024_k512_d16", + "params": { + "B": 4, + "D": 16, + "K": 512, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1836, + "evolution_flashlib_ms": 0.365599, + "evolution_kernel_ms": 0.141568, + "evolution_speedup": 2.5825, + "evolution_tflops": 0.474, + "route": "microdim_hybrid_9c0d_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d16_large_b8_n32768_k512_d16", + "params": { + "B": 8, + "D": 16, + "K": 512, + "N": 32768, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 13.0308, + "evolution_flashlib_ms": 0.3296, + "evolution_kernel_ms": 0.234416, + "evolution_speedup": 1.406, + "evolution_tflops": 18.322, + "route": "microdim_hybrid_9c0d_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d32_hugek_b1_n512_k8192_d32", + "params": { + "B": 1, + "D": 32, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1923, + "evolution_flashlib_ms": 1.396063, + "evolution_kernel_ms": 0.224064, + "evolution_speedup": 6.2306, + "evolution_tflops": 1.198, + "route": "microdim_hybrid_9c0d_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d32_paired_b2_n4096_k1024_d32", + "params": { + "B": 2, + "D": 32, + "K": 1024, + "N": 4096, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.6899, + "evolution_flashlib_ms": 0.317696, + "evolution_kernel_ms": 0.176672, + "evolution_speedup": 1.7982, + "evolution_tflops": 3.0388, + "route": "microdim_hybrid_9c0d_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d32_fallback_b3_n2432_k512_d32", + "params": { + "B": 3, + "D": 32, + "K": 512, + "N": 2432, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.8475, + "evolution_flashlib_ms": 0.28208, + "evolution_kernel_ms": 0.133824, + "evolution_speedup": 2.1078, + "evolution_tflops": 1.7865, + "route": "microdim_hybrid_9c0d_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d32_small_b4_n1024_k512_d32", + "params": { + "B": 4, + "D": 32, + "K": 512, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.3716, + "evolution_flashlib_ms": 0.361167, + "evolution_kernel_ms": 0.139583, + "evolution_speedup": 2.5875, + "evolution_tflops": 0.9616, + "route": "microdim_hybrid_9c0d_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d64_paired_b2_n4096_k1024_d64", + "params": { + "B": 2, + "D": 64, + "K": 1024, + "N": 4096, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 3.3518, + "evolution_flashlib_ms": 0.320352, + "evolution_kernel_ms": 0.070272, + "evolution_speedup": 4.5587, + "evolution_tflops": 15.2798, + "route": "d64_direct_single64_1p2gap_9f2a_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d64_fallback_b3_n2432_k512_d64", + "params": { + "B": 3, + "D": 64, + "K": 512, + "N": 2432, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.6742, + "evolution_flashlib_ms": 0.2856, + "evolution_kernel_ms": 0.067552, + "evolution_speedup": 4.2279, + "evolution_tflops": 7.0783, + "route": "d64_direct_single64_1p2gap_9f2a_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d64_small_b4_n1024_k512_d64", + "params": { + "B": 4, + "D": 64, + "K": 512, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.6741, + "evolution_flashlib_ms": 0.398223, + "evolution_kernel_ms": 0.06736, + "evolution_speedup": 5.9119, + "evolution_tflops": 3.9851, + "route": "d64_direct_single64_1p2gap_9f2a_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d64_large_b8_n32768_k512_d64", + "params": { + "B": 8, + "D": 64, + "K": 512, + "N": 32768, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 34.8042, + "evolution_flashlib_ms": 0.493615, + "evolution_kernel_ms": 0.128672, + "evolution_speedup": 3.8362, + "evolution_tflops": 133.5168, + "route": "d64_direct_single64_1p2gap_9f2a_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d80_paired_b1_n2048_k1024_d80", + "params": { + "B": 1, + "D": 80, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.8736, + "evolution_flashlib_ms": 0.384112, + "evolution_kernel_ms": 0.176896, + "evolution_speedup": 2.1714, + "evolution_tflops": 1.8968, + "route": "lowdim_e50c_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d80_small_b2_n1024_k512_d80", + "params": { + "B": 2, + "D": 80, + "K": 512, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.4431, + "evolution_flashlib_ms": 0.378656, + "evolution_kernel_ms": 0.173552, + "evolution_speedup": 2.1818, + "evolution_tflops": 0.9667, + "route": "lowdim_e50c_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d96_paired_b1_n1536_k1024_d96", + "params": { + "B": 1, + "D": 96, + "K": 1024, + "N": 1536, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.7061, + "evolution_flashlib_ms": 0.427679, + "evolution_kernel_ms": 0.176032, + "evolution_speedup": 2.4296, + "evolution_tflops": 1.7155, + "route": "lowdim_e50c_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d96_small_b2_n896_k512_d96", + "params": { + "B": 2, + "D": 96, + "K": 512, + "N": 896, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.4465, + "evolution_flashlib_ms": 0.394576, + "evolution_kernel_ms": 0.173504, + "evolution_speedup": 2.2742, + "evolution_tflops": 1.0153, + "route": "lowdim_e50c_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d96_random_b3_n6144_k1280_d96", + "params": { + "B": 3, + "D": 96, + "K": 1280, + "N": 6144, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 10.6231, + "evolution_flashlib_ms": 0.426416, + "evolution_kernel_ms": 0.177823, + "evolution_speedup": 2.398, + "evolution_tflops": 25.4739, + "route": "lowdim_e50c_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d96_fallback_b5_n1664_k256_d96", + "params": { + "B": 5, + "D": 96, + "K": 256, + "N": 1664, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.3432, + "evolution_flashlib_ms": 0.304464, + "evolution_kernel_ms": 0.171136, + "evolution_speedup": 1.7791, + "evolution_tflops": 2.3896, + "route": "lowdim_e50c_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "boundary_b1_n128_k256_d128", + "params": { + "B": 1, + "D": 128, + "K": 256, + "N": 128, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0218, + "evolution_flashlib_ms": 0.384415, + "evolution_kernel_ms": 0.143039, + "evolution_speedup": 2.6875, + "evolution_tflops": 0.0586, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_tiny_hugek_b1_n128_k4096_d128", + "params": { + "B": 1, + "D": 128, + "K": 4096, + "N": 128, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1114, + "evolution_flashlib_ms": 1.205311, + "evolution_kernel_ms": 0.171472, + "evolution_speedup": 7.0292, + "evolution_tflops": 0.7827, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "small_b1_n256_k256_d128", + "params": { + "B": 1, + "D": 128, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0385, + "evolution_flashlib_ms": 0.435263, + "evolution_kernel_ms": 0.145312, + "evolution_speedup": 2.9954, + "evolution_tflops": 0.1155, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_min_even_b1_n256_k768_d128", + "params": { + "B": 1, + "D": 128, + "K": 768, + "N": 256, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.084, + "evolution_flashlib_ms": 0.599455, + "evolution_kernel_ms": 0.167616, + "evolution_speedup": 3.5764, + "evolution_tflops": 0.3003, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_low_n_hugek_b1_n256_k4096_d128", + "params": { + "B": 1, + "D": 128, + "K": 4096, + "N": 256, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1873, + "evolution_flashlib_ms": 1.432831, + "evolution_kernel_ms": 0.209535, + "evolution_speedup": 6.8381, + "evolution_tflops": 1.2811, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "small_b1_n896_k512_d128", + "params": { + "B": 1, + "D": 128, + "K": 512, + "N": 896, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.3102, + "evolution_flashlib_ms": 0.378655, + "evolution_kernel_ms": 0.144463, + "evolution_speedup": 2.6211, + "evolution_tflops": 0.8129, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_odd_kover_b1_n896_k768_d128", + "params": { + "B": 1, + "D": 128, + "K": 768, + "N": 896, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.3949, + "evolution_flashlib_ms": 0.446111, + "evolution_kernel_ms": 0.145856, + "evolution_speedup": 3.0586, + "evolution_tflops": 1.2078, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_hugek_b1_n896_k16384_d128", + "params": { + "B": 1, + "D": 128, + "K": 16384, + "N": 896, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.6987, + "evolution_flashlib_ms": 2.21235, + "evolution_kernel_ms": 0.25696, + "evolution_speedup": 8.6097, + "evolution_tflops": 14.6252, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "boundary_b1_n1024_k768_d128", + "params": { + "B": 1, + "D": 128, + "K": 768, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.4318, + "evolution_flashlib_ms": 0.466272, + "evolution_kernel_ms": 0.169952, + "evolution_speedup": 2.7435, + "evolution_tflops": 1.1846, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_kover_b1_n1024_k2048_d128", + "params": { + "B": 1, + "D": 128, + "K": 2048, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.8902, + "evolution_flashlib_ms": 0.603104, + "evolution_kernel_ms": 0.1848, + "evolution_speedup": 3.2635, + "evolution_tflops": 2.9051, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "forced_fallback_b1_n1152_k256_d128", + "params": { + "B": 1, + "D": 128, + "K": 256, + "N": 1152, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.2411, + "evolution_flashlib_ms": 0.313135, + "evolution_kernel_ms": 0.143648, + "evolution_speedup": 2.1799, + "evolution_tflops": 0.5256, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_first_after_small_b1_n1152_k512_d128", + "params": { + "B": 1, + "D": 128, + "K": 512, + "N": 1152, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.3964, + "evolution_flashlib_ms": 0.38096, + "evolution_kernel_ms": 0.144832, + "evolution_speedup": 2.6304, + "evolution_tflops": 1.0426, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "boundary_b1_n1280_k256_d128", + "params": { + "B": 1, + "D": 128, + "K": 256, + "N": 1280, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.2722, + "evolution_flashlib_ms": 0.308176, + "evolution_kernel_ms": 0.163984, + "evolution_speedup": 1.8793, + "evolution_tflops": 0.5116, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_even_b1_n1536_k256_d128", + "params": { + "B": 1, + "D": 128, + "K": 256, + "N": 1536, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.3641, + "evolution_flashlib_ms": 0.276479, + "evolution_kernel_ms": 0.160624, + "evolution_speedup": 1.7213, + "evolution_tflops": 0.6267, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_largek_b1_n2048_k4096_d128", + "params": { + "B": 1, + "D": 128, + "K": 4096, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.9421, + "evolution_flashlib_ms": 0.729904, + "evolution_kernel_ms": 0.209024, + "evolution_speedup": 3.492, + "evolution_tflops": 10.2739, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_odd_b1_n2176_k4096_d128", + "params": { + "B": 1, + "D": 128, + "K": 4096, + "N": 2176, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 3.1992, + "evolution_flashlib_ms": 0.713215, + "evolution_kernel_ms": 0.169311, + "evolution_speedup": 4.2125, + "evolution_tflops": 13.4764, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "tail_odd_b1_n4224_k1024_d128", + "params": { + "B": 1, + "D": 128, + "K": 1024, + "N": 4224, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.9054, + "evolution_flashlib_ms": 0.38112, + "evolution_kernel_ms": 0.14784, + "evolution_speedup": 2.5779, + "evolution_tflops": 7.4898, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_large_odd_b1_n8320_k1280_d128", + "params": { + "B": 1, + "D": 128, + "K": 1280, + "N": 8320, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 7.8828, + "evolution_flashlib_ms": 0.345856, + "evolution_kernel_ms": 0.149984, + "evolution_speedup": 2.306, + "evolution_tflops": 18.1773, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_large_b1_n65536_k4096_d128", + "params": { + "B": 1, + "D": 128, + "K": 4096, + "N": 65536, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 62.6692, + "evolution_flashlib_ms": 1.096543, + "evolution_kernel_ms": 0.259808, + "evolution_speedup": 4.2206, + "evolution_tflops": 264.501, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_large_n_b1_n131072_k256_d128", + "params": { + "B": 1, + "D": 128, + "K": 256, + "N": 131072, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 31.6907, + "evolution_flashlib_ms": 0.271056, + "evolution_kernel_ms": 0.184608, + "evolution_speedup": 1.4683, + "evolution_tflops": 46.5307, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_kover_b2_n256_k768_d128", + "params": { + "B": 2, + "D": 128, + "K": 768, + "N": 256, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1597, + "evolution_flashlib_ms": 0.630431, + "evolution_kernel_ms": 0.16816, + "evolution_speedup": 3.749, + "evolution_tflops": 0.5986, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_low_n_hugek_b2_n384_k8192_d128", + "params": { + "B": 2, + "D": 128, + "K": 8192, + "N": 384, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.841, + "evolution_flashlib_ms": 1.915198, + "evolution_kernel_ms": 0.20048, + "evolution_speedup": 9.5531, + "evolution_tflops": 8.0338, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_even_b2_n512_k8192_d128", + "params": { + "B": 2, + "D": 128, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.2615, + "evolution_flashlib_ms": 1.702303, + "evolution_kernel_ms": 0.269792, + "evolution_speedup": 6.3097, + "evolution_tflops": 7.9598, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "small_random_b2_n768_k512_d128", + "params": { + "B": 2, + "D": 128, + "K": 512, + "N": 768, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.4722, + "evolution_flashlib_ms": 0.426368, + "evolution_kernel_ms": 0.144992, + "evolution_speedup": 2.9406, + "evolution_tflops": 1.3885, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "random_b2_n768_k3072_d128", + "params": { + "B": 2, + "D": 128, + "K": 3072, + "N": 768, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.4768, + "evolution_flashlib_ms": 0.817951, + "evolution_kernel_ms": 0.319424, + "evolution_speedup": 2.5607, + "evolution_tflops": 3.7817, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "smoke_b2_n1024_k512_d128", + "params": { + "B": 2, + "D": 128, + "K": 512, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.6361, + "evolution_flashlib_ms": 0.422031, + "evolution_kernel_ms": 0.148896, + "evolution_speedup": 2.8344, + "evolution_tflops": 1.8028, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_first_odd_b2_n1408_k256_d128", + "params": { + "B": 2, + "D": 128, + "K": 256, + "N": 1408, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.6573, + "evolution_flashlib_ms": 0.280752, + "evolution_kernel_ms": 0.14272, + "evolution_speedup": 1.9672, + "evolution_tflops": 1.2931, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_even_tail_b2_n1792_k256_d128", + "params": { + "B": 2, + "D": 128, + "K": 256, + "N": 1792, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.9451, + "evolution_flashlib_ms": 0.248512, + "evolution_kernel_ms": 0.164208, + "evolution_speedup": 1.5134, + "evolution_tflops": 1.4304, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "tail_even_b2_n2048_k2048_d128", + "params": { + "B": 2, + "D": 128, + "K": 2048, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 3.8606, + "evolution_flashlib_ms": 0.556255, + "evolution_kernel_ms": 0.186303, + "evolution_speedup": 2.9858, + "evolution_tflops": 11.5268, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "random_legal_b2_n4736_k512_d128", + "params": { + "B": 2, + "D": 128, + "K": 512, + "N": 4736, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 3.587, + "evolution_flashlib_ms": 0.346112, + "evolution_kernel_ms": 0.146688, + "evolution_speedup": 2.3595, + "evolution_tflops": 8.4636, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_large_b2_n65664_k512_d128", + "params": { + "B": 2, + "D": 128, + "K": 512, + "N": 65664, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 40.2273, + "evolution_flashlib_ms": 0.427904, + "evolution_kernel_ms": 0.17792, + "evolution_speedup": 2.405, + "evolution_tflops": 96.7478, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "small_low_b3_n256_k512_d128", + "params": { + "B": 3, + "D": 128, + "K": 512, + "N": 256, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1812, + "evolution_flashlib_ms": 0.555391, + "evolution_kernel_ms": 0.145791, + "evolution_speedup": 3.8095, + "evolution_tflops": 0.6905, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_first_even_b3_n1536_k512_d128", + "params": { + "B": 3, + "D": 128, + "K": 512, + "N": 1536, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.7141, + "evolution_flashlib_ms": 0.352351, + "evolution_kernel_ms": 0.166016, + "evolution_speedup": 2.1224, + "evolution_tflops": 3.6381, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_mid_odd_b3_n2432_k512_d128", + "params": { + "B": 3, + "D": 128, + "K": 512, + "N": 2432, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 3.0254, + "evolution_flashlib_ms": 0.316095, + "evolution_kernel_ms": 0.144, + "evolution_speedup": 2.1951, + "evolution_tflops": 6.641, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "random_legal_b3_n2560_k1536_d128", + "params": { + "B": 3, + "D": 128, + "K": 1536, + "N": 2560, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 7.0022, + "evolution_flashlib_ms": 0.43128, + "evolution_kernel_ms": 0.179008, + "evolution_speedup": 2.4093, + "evolution_tflops": 16.8702, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "random_b3_n6144_k1280_d128", + "params": { + "B": 3, + "D": 128, + "K": 1280, + "N": 6144, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 14.1328, + "evolution_flashlib_ms": 0.427359, + "evolution_kernel_ms": 0.17632, + "evolution_speedup": 2.4238, + "evolution_tflops": 34.2548, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "small_b4_n128_k512_d128", + "params": { + "B": 4, + "D": 128, + "K": 512, + "N": 128, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1386, + "evolution_flashlib_ms": 0.484159, + "evolution_kernel_ms": 0.144912, + "evolution_speedup": 3.3411, + "evolution_tflops": 0.4631, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "random_legal_b4_n640_k1024_d128", + "params": { + "B": 4, + "D": 128, + "K": 1024, + "N": 640, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.2317, + "evolution_flashlib_ms": 0.544863, + "evolution_kernel_ms": 0.148095, + "evolution_speedup": 3.6791, + "evolution_tflops": 4.5315, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "boundary_b4_n1024_k512_d128", + "params": { + "B": 4, + "D": 128, + "K": 512, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.3605, + "evolution_flashlib_ms": 0.394623, + "evolution_kernel_ms": 0.144512, + "evolution_speedup": 2.7307, + "evolution_tflops": 3.7151, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_odd_b4_n1664_k256_d128", + "params": { + "B": 4, + "D": 128, + "K": 256, + "N": 1664, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.4791, + "evolution_flashlib_ms": 0.294912, + "evolution_kernel_ms": 0.145728, + "evolution_speedup": 2.0237, + "evolution_tflops": 2.9933, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_widek_b4_n3072_k3072_d128", + "params": { + "B": 4, + "D": 128, + "K": 3072, + "N": 3072, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 15.4254, + "evolution_flashlib_ms": 0.62648, + "evolution_kernel_ms": 0.19712, + "evolution_speedup": 3.1782, + "evolution_tflops": 49.0242, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "random_legal_b5_n512_k512_d128", + "params": { + "B": 5, + "D": 128, + "K": 512, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.7414, + "evolution_flashlib_ms": 0.452576, + "evolution_kernel_ms": 0.14368, + "evolution_speedup": 3.1499, + "evolution_tflops": 2.3354, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "random_b5_n3456_k512_d128", + "params": { + "B": 5, + "D": 128, + "K": 512, + "N": 3456, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 7.5675, + "evolution_flashlib_ms": 0.299296, + "evolution_kernel_ms": 0.145872, + "evolution_speedup": 2.0518, + "evolution_tflops": 15.5268, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "small_cap_b6_n1024_k256_d128", + "params": { + "B": 6, + "D": 128, + "K": 256, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.141, + "evolution_flashlib_ms": 0.352896, + "evolution_kernel_ms": 0.144144, + "evolution_speedup": 2.4482, + "evolution_tflops": 2.7934, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_large_b6_n65536_k512_d128", + "params": { + "B": 6, + "D": 128, + "K": 512, + "N": 65536, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 57.2499, + "evolution_flashlib_ms": 0.900256, + "evolution_kernel_ms": 0.253025, + "evolution_speedup": 3.558, + "evolution_tflops": 203.6937, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "small_boundary_b7_n896_k256_d128", + "params": { + "B": 7, + "D": 128, + "K": 256, + "N": 896, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.2543, + "evolution_flashlib_ms": 0.327696, + "evolution_kernel_ms": 0.14368, + "evolution_speedup": 2.2807, + "evolution_tflops": 2.8608, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "random_b7_n1024_k4096_d128", + "params": { + "B": 7, + "D": 128, + "K": 4096, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 8.1063, + "evolution_flashlib_ms": 0.9272, + "evolution_kernel_ms": 0.20928, + "evolution_speedup": 4.4304, + "evolution_tflops": 35.9145, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "fallback_b8_n3712_k2048_d128", + "params": { + "B": 8, + "D": 128, + "K": 2048, + "N": 3712, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 25.0897, + "evolution_flashlib_ms": 0.620543, + "evolution_kernel_ms": 0.171536, + "evolution_speedup": 3.6176, + "evolution_tflops": 90.764, + "route": "aligned_v10_fallback", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paired_b8_n4096_k256_d128", + "params": { + "B": 8, + "D": 128, + "K": 256, + "N": 4096, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 8.9324, + "evolution_flashlib_ms": 0.240416, + "evolution_kernel_ms": 0.164624, + "evolution_speedup": 1.4604, + "evolution_tflops": 13.0448, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "mid_b8_n8192_k512_d128", + "params": { + "B": 8, + "D": 128, + "K": 512, + "N": 8192, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 21.7357, + "evolution_flashlib_ms": 0.395199, + "evolution_kernel_ms": 0.176959, + "evolution_speedup": 2.2333, + "evolution_tflops": 48.5418, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "random_legal_b9_n3840_k1792_d128", + "params": { + "B": 9, + "D": 128, + "K": 1792, + "N": 3840, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 26.2993, + "evolution_flashlib_ms": 0.602848, + "evolution_kernel_ms": 0.186464, + "evolution_speedup": 3.2331, + "evolution_tflops": 85.027, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "small_b12_n640_k256_d128", + "params": { + "B": 12, + "D": 128, + "K": 256, + "N": 640, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 1.3627, + "evolution_flashlib_ms": 0.36936, + "evolution_kernel_ms": 0.144128, + "evolution_speedup": 2.5627, + "evolution_tflops": 3.4921, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "small_b16_n1024_k512_d128", + "params": { + "B": 16, + "D": 128, + "K": 512, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 5.3198, + "evolution_flashlib_ms": 0.40368, + "evolution_kernel_ms": 0.1464, + "evolution_speedup": 2.7574, + "evolution_tflops": 14.6686, + "route": "small_grid_single_tile_v10", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "large_b16_n32768_k1024_d128", + "params": { + "B": 16, + "D": 128, + "K": 1024, + "N": 32768, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 65.9244, + "evolution_flashlib_ms": 2.084797, + "evolution_kernel_ms": 0.363695, + "evolution_speedup": 5.7323, + "evolution_tflops": 377.8962, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "paper_b32_n75776_k1024_d128", + "params": { + "B": 32, + "D": 128, + "K": 1024, + "N": 75776, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 75.3887, + "evolution_flashlib_ms": 8.431702, + "evolution_kernel_ms": 0.951326, + "evolution_speedup": 8.8631, + "evolution_tflops": 668.1781, + "route": "paired_large_v15", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d160_paired_b1_n2048_k2048_d160", + "params": { + "B": 1, + "D": 160, + "K": 2048, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.3001, + "evolution_flashlib_ms": 0.583535, + "evolution_kernel_ms": 0.188448, + "evolution_speedup": 3.0965, + "evolution_tflops": 7.1223, + "route": "d160_padded_single_repeated_mma_v2", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d160_fallback_b2_n2432_k1024_d160", + "params": { + "B": 2, + "D": 160, + "K": 1024, + "N": 2432, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 3.798, + "evolution_flashlib_ms": 0.419648, + "evolution_kernel_ms": 0.18832, + "evolution_speedup": 2.2284, + "evolution_tflops": 8.4634, + "route": "d160_padded_single_repeated_mma_v2", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d192_fallback_b1_n2176_k1024_d192", + "params": { + "B": 1, + "D": 192, + "K": 1024, + "N": 2176, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 2.0726, + "evolution_flashlib_ms": 0.412832, + "evolution_kernel_ms": 0.149952, + "evolution_speedup": 2.7531, + "evolution_tflops": 5.7061, + "route": "d192_single_repeated_mma_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d192_small_b2_n1024_k512_d192", + "params": { + "B": 2, + "D": 192, + "K": 512, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.8948, + "evolution_flashlib_ms": 0.449999, + "evolution_kernel_ms": 0.146657, + "evolution_speedup": 3.0684, + "evolution_tflops": 2.7455, + "route": "d192_single_repeated_mma_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d192_paired_b2_n2048_k2048_d192", + "params": { + "B": 2, + "D": 192, + "K": 2048, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 4.6837, + "evolution_flashlib_ms": 0.687759, + "evolution_kernel_ms": 0.310336, + "evolution_speedup": 2.2162, + "evolution_tflops": 10.3798, + "route": "d192_paired_repeated_mma_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d192_large_b4_n32768_k1024_d192", + "params": { + "B": 4, + "D": 192, + "K": 1024, + "N": 32768, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 56.4296, + "evolution_flashlib_ms": 0.913344, + "evolution_kernel_ms": 0.221536, + "evolution_speedup": 4.1228, + "evolution_tflops": 232.6466, + "route": "d192_paired_repeated_mma_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d256_hugek_b1_n512_k8192_d256", + "params": { + "B": 1, + "D": 256, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0389, + "evolution_flashlib_ms": 55.197531, + "evolution_kernel_ms": 0.222944, + "evolution_speedup": 247.5847, + "evolution_tflops": 9.6324, + "route": "d256_single_repeated_mma_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d256_small_b1_n1024_k512_d256", + "params": { + "B": 1, + "D": 256, + "K": 512, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.6249, + "evolution_flashlib_ms": 0.429567, + "evolution_kernel_ms": 0.148224, + "evolution_speedup": 2.8981, + "evolution_tflops": 1.811, + "route": "d256_single_repeated_mma_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d256_paired_b1_n4096_k4096_d256", + "params": { + "B": 1, + "D": 256, + "K": 4096, + "N": 4096, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1557, + "evolution_flashlib_ms": 55.152968, + "evolution_kernel_ms": 0.179455, + "evolution_speedup": 307.3359, + "evolution_tflops": 47.8668, + "route": "d256_single_repeated_mma_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d256_fallback_b2_n2432_k2048_d256", + "params": { + "B": 2, + "D": 256, + "K": 2048, + "N": 2432, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 8.2786, + "evolution_flashlib_ms": 0.616079, + "evolution_kernel_ms": 0.16352, + "evolution_speedup": 3.7676, + "evolution_tflops": 31.1905, + "route": "d256_single_repeated_mma_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d320_paired_b1_n2048_k4096_d320", + "params": { + "B": 1, + "D": 320, + "K": 4096, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1154, + "evolution_flashlib_ms": 46.53602, + "evolution_kernel_ms": 0.23344, + "evolution_speedup": 199.349, + "evolution_tflops": 22.9982, + "route": "highd_splitd_single_tile_6fcf_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d320_large_b2_n16384_k1024_d320", + "params": { + "B": 2, + "D": 320, + "K": 1024, + "N": 16384, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 36.5995, + "evolution_flashlib_ms": 0.586752, + "evolution_kernel_ms": 0.19096, + "evolution_speedup": 3.0726, + "evolution_tflops": 112.4573, + "route": "highd_splitd_single_tile_6fcf_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d384_hugek_b1_n768_k8192_d384", + "params": { + "B": 1, + "D": 384, + "K": 8192, + "N": 768, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1211, + "evolution_flashlib_ms": 39.895303, + "evolution_kernel_ms": 0.271664, + "evolution_speedup": 146.8553, + "evolution_tflops": 17.7861, + "route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d384_small_b1_n896_k512_d384", + "params": { + "B": 1, + "D": 384, + "K": 512, + "N": 896, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.6757, + "evolution_flashlib_ms": 0.521424, + "evolution_kernel_ms": 0.15824, + "evolution_speedup": 3.2951, + "evolution_tflops": 2.2265, + "route": "highd_splitd_single_tile_6fcf_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d384_paired_b1_n2048_k4096_d384", + "params": { + "B": 1, + "D": 384, + "K": 4096, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1674, + "evolution_flashlib_ms": 38.477384, + "evolution_kernel_ms": 0.252672, + "evolution_speedup": 152.2819, + "evolution_tflops": 25.4973, + "route": "highd_splitd_single_tile_6fcf_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d384_fallback_b3_n3456_k1024_d384", + "params": { + "B": 3, + "D": 384, + "K": 1024, + "N": 3456, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 13.6962, + "evolution_flashlib_ms": 0.595328, + "evolution_kernel_ms": 0.174575, + "evolution_speedup": 3.4101, + "evolution_tflops": 46.706, + "route": "highd_splitd_single_tile_6fcf_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d448_hugek_b1_n512_k8192_d448", + "params": { + "B": 1, + "D": 448, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.0923, + "evolution_flashlib_ms": 40.698431, + "evolution_kernel_ms": 0.262959, + "evolution_speedup": 154.771, + "evolution_tflops": 14.2916, + "route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d448_paired_b1_n2048_k4096_d448", + "params": { + "B": 1, + "D": 448, + "K": 4096, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1943, + "evolution_flashlib_ms": 38.67485, + "evolution_kernel_ms": 0.269632, + "evolution_speedup": 143.4357, + "evolution_tflops": 27.8757, + "route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d512_hugek_b1_n512_k8192_d512", + "params": { + "B": 1, + "D": 512, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.1326, + "evolution_flashlib_ms": 32.378557, + "evolution_kernel_ms": 0.266288, + "evolution_speedup": 121.5925, + "evolution_tflops": 16.1291, + "route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d512_small_b1_n1024_k512_d512", + "params": { + "B": 1, + "D": 512, + "K": 512, + "N": 1024, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.969, + "evolution_flashlib_ms": 0.554048, + "evolution_kernel_ms": 0.162688, + "evolution_speedup": 3.4056, + "evolution_tflops": 3.3, + "route": "highd_splitd_single_tile_6fcf_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d512_paired_b1_n2048_k4096_d512", + "params": { + "B": 1, + "D": 512, + "K": 4096, + "N": 2048, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 0.281, + "evolution_flashlib_ms": 30.573365, + "evolution_kernel_ms": 0.270976, + "evolution_speedup": 112.8268, + "evolution_tflops": 31.7, + "route": "flash_kmeans_assign_no_padding_portfolio_r63_g1_d512n512_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "d512_fallback_b2_n2432_k2048_d512", + "params": { + "B": 2, + "D": 512, + "K": 2048, + "N": 2432, + "dtype": "bfloat16", + "evolution_flashlib_equiv_tflops": 11.9237, + "evolution_flashlib_ms": 0.855488, + "evolution_kernel_ms": 0.220319, + "evolution_speedup": 3.883, + "evolution_tflops": 46.299, + "route": "highd_splitd_single_tile_6fcf_v1", + "source": "full95" + }, + "recorded": {}, + "runtime_coverage": true + }, + { + "label": "post_d895_d16_b4_n32768_k1024_d16", + "params": { + "B": 4, + "D": 16, + "K": 1024, + "N": 32768, + "dtype": "bfloat16", + "seed": 21602, + "source": "near_floor_microdim" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 12.993, + "evolution_flashlib_ms": 0.330559, + "evolution_kernel_ms": 0.216591, + "evolution_route": "microdim_hybrid_9c0d_v1", + "evolution_speedup": 1.5262, + "evolution_tflops": 19.8298 + } + }, + { + "label": "post_d895_d16_b8_n65536_k512_d16", + "params": { + "B": 8, + "D": 16, + "K": 512, + "N": 65536, + "dtype": "bfloat16", + "seed": 21601, + "source": "near_floor_microdim" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 19.3119, + "evolution_flashlib_ms": 0.444799, + "evolution_kernel_ms": 0.295359, + "evolution_route": "microdim_hybrid_9c0d_v1", + "evolution_speedup": 1.506, + "evolution_tflops": 29.083 + } + }, + { + "label": "post_d895_d32_b4_n32768_k1024_d32", + "params": { + "B": 4, + "D": 32, + "K": 1024, + "N": 32768, + "dtype": "bfloat16", + "seed": 23202, + "source": "near_floor_microdim" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 23.8017, + "evolution_flashlib_ms": 0.360896, + "evolution_kernel_ms": 0.218816, + "evolution_route": "microdim_hybrid_9c0d_v1", + "evolution_speedup": 1.6493, + "evolution_tflops": 39.2564 + } + }, + { + "label": "post_d895_d32_b8_n65536_k512_d32", + "params": { + "B": 8, + "D": 32, + "K": 512, + "N": 65536, + "dtype": "bfloat16", + "seed": 23201, + "source": "near_floor_microdim" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 34.3587, + "evolution_flashlib_ms": 0.500015, + "evolution_kernel_ms": 0.301439, + "evolution_route": "microdim_hybrid_9c0d_v1", + "evolution_speedup": 1.6588, + "evolution_tflops": 56.9928 + } + }, + { + "label": "adjacent_3328_d48_small_boundary_b1_n256_k256_d48", + "params": { + "B": 1, + "D": 48, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "seed": 3328481, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0145, + "evolution_flashlib_ms": 0.434063, + "evolution_kernel_ms": 0.168656, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.5737, + "evolution_tflops": 0.0373 + } + }, + { + "label": "post_d895_d48_b1_n512_k8192_d48", + "params": { + "B": 1, + "D": 48, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 24803, + "source": "high_k_low_n" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.275, + "evolution_flashlib_ms": 1.464301, + "evolution_kernel_ms": 0.21952, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 6.6705, + "evolution_tflops": 1.8342 + } + }, + { + "label": "adjacent_d9d5_d48_highk_heldout_b2_n640_k4096_d48", + "params": { + "B": 2, + "D": 48, + "K": 4096, + "N": 640, + "dtype": "bfloat16", + "seed": 9504801, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.5494, + "evolution_flashlib_ms": 0.916111, + "evolution_kernel_ms": 0.19392, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 4.7242, + "evolution_tflops": 2.5955 + } + }, + { + "label": "adjacent_a2f8_d48_k1024_boundary_b2_n768_k1024_d48", + "params": { + "B": 2, + "D": 48, + "K": 1024, + "N": 768, + "dtype": "bfloat16", + "seed": 1028481, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.3038, + "evolution_flashlib_ms": 0.497055, + "evolution_kernel_ms": 0.173184, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.8701, + "evolution_tflops": 0.8719 + } + }, + { + "label": "adjacent_8f09_d48_medium_tail_b2_n1792_k512_d48", + "params": { + "B": 2, + "D": 48, + "K": 512, + "N": 1792, + "dtype": "bfloat16", + "seed": 8090481, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.556, + "evolution_flashlib_ms": 0.316848, + "evolution_kernel_ms": 0.169375, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.8707, + "evolution_tflops": 1.0401 + } + }, + { + "label": "post_d895_d48_b2_n2048_k512_d48", + "params": { + "B": 2, + "D": 48, + "K": 512, + "N": 2048, + "dtype": "bfloat16", + "seed": 24801, + "source": "new_lowmid_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.6574, + "evolution_flashlib_ms": 0.30624, + "evolution_kernel_ms": 0.166656, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.8376, + "evolution_tflops": 1.208 + } + }, + { + "label": "adjacent_5600_d48_low_n_boundary_b3_n1536_k256_d48", + "params": { + "B": 3, + "D": 48, + "K": 256, + "N": 1536, + "dtype": "bfloat16", + "seed": 5600481, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.4029, + "evolution_flashlib_ms": 0.281087, + "evolution_kernel_ms": 0.169183, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.6614, + "evolution_tflops": 0.6694 + } + }, + { + "label": "adjacent_c44f_d48_midk_boundary_b3_n2688_k768_d48", + "params": { + "B": 3, + "D": 48, + "K": 768, + "N": 2688, + "dtype": "bfloat16", + "seed": 4404801, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 1.9182, + "evolution_flashlib_ms": 0.309952, + "evolution_kernel_ms": 0.174383, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.7774, + "evolution_tflops": 3.4094 + } + }, + { + "label": "adjacent_ef00_d48_midn_boundary_b3_n3456_k512_d48", + "params": { + "B": 3, + "D": 48, + "K": 512, + "N": 3456, + "dtype": "bfloat16", + "seed": 9000481, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 1.9767, + "evolution_flashlib_ms": 0.257807, + "evolution_kernel_ms": 0.172495, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.4946, + "evolution_tflops": 2.9543 + } + }, + { + "label": "adjacent_68cf_d48_boundary_b4_n2304_k512_d48", + "params": { + "B": 4, + "D": 48, + "K": 512, + "N": 2304, + "dtype": "bfloat16", + "seed": 6804801, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 1.5408, + "evolution_flashlib_ms": 0.293984, + "evolution_kernel_ms": 0.170608, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.7232, + "evolution_tflops": 2.6551 + } + }, + { + "label": "adjacent_1d49_d48_lowk_boundary_b4_n3968_k256_d48", + "params": { + "B": 4, + "D": 48, + "K": 256, + "N": 3968, + "dtype": "bfloat16", + "seed": 1490481, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 1.7051, + "evolution_flashlib_ms": 0.228768, + "evolution_kernel_ms": 0.17192, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.3307, + "evolution_tflops": 2.2689 + } + }, + { + "label": "post_d895_d48_b4_n8192_k1024_d48", + "params": { + "B": 4, + "D": 48, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "seed": 24802, + "source": "new_lowmid_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 9.5879, + "evolution_flashlib_ms": 0.335967, + "evolution_kernel_ms": 0.179807, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.8685, + "evolution_tflops": 17.9149 + } + }, + { + "label": "adjacent_9ca0_d48_guard_boundary_b6_n12288_k512_d48", + "params": { + "B": 6, + "D": 48, + "K": 512, + "N": 12288, + "dtype": "bfloat16", + "seed": 924805, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 11.385, + "evolution_flashlib_ms": 0.318304, + "evolution_kernel_ms": 0.182336, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.7457, + "evolution_tflops": 19.8747 + } + }, + { + "label": "post_d895_d112_b1_n256_k256_d112", + "params": { + "B": 1, + "D": 112, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "seed": 211204, + "source": "lowmid_small_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0353, + "evolution_flashlib_ms": 0.416031, + "evolution_kernel_ms": 0.166575, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.4976, + "evolution_tflops": 0.0881 + } + }, + { + "label": "adjacent_8f09_d112_small_highk_b1_n384_k4096_d112", + "params": { + "B": 1, + "D": 112, + "K": 4096, + "N": 384, + "dtype": "bfloat16", + "seed": 8091121, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.2782, + "evolution_flashlib_ms": 1.266206, + "evolution_kernel_ms": 0.195296, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 6.4835, + "evolution_tflops": 1.804 + } + }, + { + "label": "post_d895_d112_b1_n512_k8192_d112", + "params": { + "B": 1, + "D": 112, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 211203, + "source": "high_k_low_n" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.5425, + "evolution_flashlib_ms": 1.731693, + "evolution_kernel_ms": 0.294304, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 5.884, + "evolution_tflops": 3.1924 + } + }, + { + "label": "adjacent_1d49_d112_highk_tail_b1_n1408_k4096_d112", + "params": { + "B": 1, + "D": 112, + "K": 4096, + "N": 1408, + "dtype": "bfloat16", + "seed": 1491121, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 1.4908, + "evolution_flashlib_ms": 0.866559, + "evolution_kernel_ms": 0.196192, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 4.4169, + "evolution_tflops": 6.5846 + } + }, + { + "label": "adjacent_3328_d112_highk_low_n_b2_n512_k8192_d112", + "params": { + "B": 2, + "D": 112, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 3328112, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 1.0686, + "evolution_flashlib_ms": 1.758429, + "evolution_kernel_ms": 0.295872, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 5.9432, + "evolution_tflops": 6.3509 + } + }, + { + "label": "post_d895_d112_b2_n2048_k512_d112", + "params": { + "B": 2, + "D": 112, + "K": 512, + "N": 2048, + "dtype": "bfloat16", + "seed": 211201, + "source": "new_lowmid_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 1.5026, + "evolution_flashlib_ms": 0.312623, + "evolution_kernel_ms": 0.188, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.6629, + "evolution_tflops": 2.4987 + } + }, + { + "label": "adjacent_c44f_d112_odd_tail_b2_n3200_k768_d112", + "params": { + "B": 2, + "D": 112, + "K": 768, + "N": 3200, + "dtype": "bfloat16", + "seed": 4411201, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 3.4225, + "evolution_flashlib_ms": 0.321696, + "evolution_kernel_ms": 0.173088, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.8586, + "evolution_tflops": 6.361 + } + }, + { + "label": "adjacent_d9d5_d112_highk_request_b3_n768_k8192_d112", + "params": { + "B": 3, + "D": 112, + "K": 8192, + "N": 768, + "dtype": "bfloat16", + "seed": 9511201, + "source": "request_specific" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.7846, + "evolution_flashlib_ms": 1.518318, + "evolution_kernel_ms": 0.287872, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 5.2743, + "evolution_tflops": 14.6866 + } + }, + { + "label": "adjacent_9ca0_d112_tail_div_b3_n3840_k768_d112", + "params": { + "B": 3, + "D": 112, + "K": 768, + "N": 3840, + "dtype": "bfloat16", + "seed": 921125, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 5.9472, + "evolution_flashlib_ms": 0.333232, + "evolution_kernel_ms": 0.194592, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.7125, + "evolution_tflops": 10.1844 + } + }, + { + "label": "adjacent_a2f8_d112_lowk_tail_b4_n3456_k256_d112", + "params": { + "B": 4, + "D": 112, + "K": 256, + "N": 3456, + "dtype": "bfloat16", + "seed": 1028112, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 3.4264, + "evolution_flashlib_ms": 0.231359, + "evolution_kernel_ms": 0.168912, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.3697, + "evolution_tflops": 4.6931 + } + }, + { + "label": "adjacent_ef00_d112_odd_tail_b4_n3712_k1024_d112", + "params": { + "B": 4, + "D": 112, + "K": 1024, + "N": 3712, + "dtype": "bfloat16", + "seed": 9001121, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 8.805, + "evolution_flashlib_ms": 0.386799, + "evolution_kernel_ms": 0.176608, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.1902, + "evolution_tflops": 19.2844 + } + }, + { + "label": "post_d895_d112_b4_n8192_k1024_d112", + "params": { + "B": 4, + "D": 112, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "seed": 211202, + "source": "new_lowmid_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 18.1068, + "evolution_flashlib_ms": 0.415103, + "evolution_kernel_ms": 0.1956, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.1222, + "evolution_tflops": 38.4264 + } + }, + { + "label": "adjacent_5600_d112_odd_tile_tail_b5_n2176_k512_d112", + "params": { + "B": 5, + "D": 112, + "K": 512, + "N": 2176, + "dtype": "bfloat16", + "seed": 5601121, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 3.6867, + "evolution_flashlib_ms": 0.338463, + "evolution_kernel_ms": 0.171519, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.9733, + "evolution_tflops": 7.275 + } + }, + { + "label": "adjacent_68cf_d112_tail_b5_n2944_k512_d112", + "params": { + "B": 5, + "D": 112, + "K": 512, + "N": 2944, + "dtype": "bfloat16", + "seed": 6811201, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 5.7016, + "evolution_flashlib_ms": 0.296095, + "evolution_kernel_ms": 0.173023, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.7113, + "evolution_tflops": 9.7571 + } + }, + { + "label": "post_d895_d128_paired_b2_n262144_k256_d128", + "params": { + "B": 2, + "D": 128, + "K": 256, + "N": 262144, + "dtype": "bfloat16", + "seed": 212802, + "source": "near_floor_paired_d128" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 51.8017, + "evolution_flashlib_ms": 0.663294, + "evolution_kernel_ms": 0.24464, + "evolution_route": "paired_large_v15", + "evolution_speedup": 2.7113, + "evolution_tflops": 140.4505 + } + }, + { + "label": "post_d895_d128_fallback_b3_n1920_k256_d128", + "params": { + "B": 3, + "D": 128, + "K": 256, + "N": 1920, + "dtype": "bfloat16", + "seed": 212803, + "source": "near_floor_fallback_d128" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 1.3402, + "evolution_flashlib_ms": 0.281664, + "evolution_kernel_ms": 0.1436, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 1.9614, + "evolution_tflops": 2.6287 + } + }, + { + "label": "adjacent_8f09_d128_fallback_neighbor_b4_n4480_k512_d128", + "params": { + "B": 4, + "D": 128, + "K": 512, + "N": 4480, + "dtype": "bfloat16", + "seed": 8091281, + "source": "guard_miss_fallback" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 7.2152, + "evolution_flashlib_ms": 0.325535, + "evolution_kernel_ms": 0.145856, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 2.2319, + "evolution_tflops": 16.1036 + } + }, + { + "label": "adjacent_ef00_d128_forced_fallback_b4_n7552_k512_d128", + "params": { + "B": 4, + "D": 128, + "K": 512, + "N": 7552, + "dtype": "bfloat16", + "seed": 9001281, + "source": "forced_fallback" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 12.6084, + "evolution_flashlib_ms": 0.314032, + "evolution_kernel_ms": 0.153023, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 2.0522, + "evolution_tflops": 25.8747 + } + }, + { + "label": "post_d895_d128_fallback_b5_n2176_k512_d128", + "params": { + "B": 5, + "D": 128, + "K": 512, + "N": 2176, + "dtype": "bfloat16", + "seed": 212804, + "source": "near_floor_fallback_d128" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 4.6116, + "evolution_flashlib_ms": 0.309232, + "evolution_kernel_ms": 0.14432, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 2.1427, + "evolution_tflops": 9.8813 + } + }, + { + "label": "adjacent_d9d5_d128_guard_miss_fallback_b5_n6016_k1024_d128", + "params": { + "B": 5, + "D": 128, + "K": 1024, + "N": 6016, + "dtype": "bfloat16", + "seed": 9512801, + "source": "guard_miss_fallback" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 17.1342, + "evolution_flashlib_ms": 0.460207, + "evolution_kernel_ms": 0.159104, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 2.8925, + "evolution_tflops": 49.5606 + } + }, + { + "label": "adjacent_1d49_d128_forced_fallback_b5_n7296_k512_d128", + "params": { + "B": 5, + "D": 128, + "K": 512, + "N": 7296, + "dtype": "bfloat16", + "seed": 1491281, + "source": "forced_fallback" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 14.0693, + "evolution_flashlib_ms": 0.339855, + "evolution_kernel_ms": 0.152639, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 2.2265, + "evolution_tflops": 31.3255 + } + }, + { + "label": "adjacent_c44f_d128_forced_fallback_b6_n6272_k512_d128", + "params": { + "B": 6, + "D": 128, + "K": 512, + "N": 6272, + "dtype": "bfloat16", + "seed": 4412801, + "source": "forced_fallback" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 14.7956, + "evolution_flashlib_ms": 0.333376, + "evolution_kernel_ms": 0.152256, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 2.1896, + "evolution_tflops": 32.3961 + } + }, + { + "label": "adjacent_a2f8_d128_odd_fallback_b6_n8576_k512_d128", + "params": { + "B": 6, + "D": 128, + "K": 512, + "N": 8576, + "dtype": "bfloat16", + "seed": 1028128, + "source": "guard_miss_fallback" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 19.2761, + "evolution_flashlib_ms": 0.349887, + "evolution_kernel_ms": 0.158688, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 2.2049, + "evolution_tflops": 42.5013 + } + }, + { + "label": "post_d895_d128_fallback_b7_n2432_k1024_d128", + "params": { + "B": 7, + "D": 128, + "K": 1024, + "N": 2432, + "dtype": "bfloat16", + "seed": 212805, + "source": "near_floor_fallback_d128" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 11.0028, + "evolution_flashlib_ms": 0.4056, + "evolution_kernel_ms": 0.147936, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 2.7417, + "evolution_tflops": 30.1667 + } + }, + { + "label": "adjacent_68cf_d128_forced_fallback_b7_n6016_k512_d128", + "params": { + "B": 7, + "D": 128, + "K": 512, + "N": 6016, + "dtype": "bfloat16", + "seed": 6812801, + "source": "forced_fallback" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 14.2803, + "evolution_flashlib_ms": 0.386527, + "evolution_kernel_ms": 0.156832, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 2.4646, + "evolution_tflops": 35.195 + } + }, + { + "label": "adjacent_9ca0_d128_exact_guard_neighbor_b8_n8064_k256_d128", + "params": { + "B": 8, + "D": 128, + "K": 256, + "N": 8064, + "dtype": "bfloat16", + "seed": 912806, + "source": "guard_miss_fallback" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 14.5156, + "evolution_flashlib_ms": 0.291264, + "evolution_kernel_ms": 0.155168, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 1.8771, + "evolution_tflops": 27.247 + } + }, + { + "label": "post_d895_d128_paired_b8_n8192_k256_d128", + "params": { + "B": 8, + "D": 128, + "K": 256, + "N": 8192, + "dtype": "bfloat16", + "seed": 212801, + "source": "near_floor_paired_d128" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 14.6751, + "evolution_flashlib_ms": 0.292671, + "evolution_kernel_ms": 0.15424, + "evolution_route": "d128_even_near_floor_v10_repair", + "evolution_speedup": 1.8975, + "evolution_tflops": 27.846 + } + }, + { + "label": "adjacent_3328_d128_exact_guard_k_neighbor_b8_n8192_k512_d128", + "params": { + "B": 8, + "D": 128, + "K": 512, + "N": 8192, + "dtype": "bfloat16", + "seed": 3328128, + "source": "guard_miss_fallback" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 21.7692, + "evolution_flashlib_ms": 0.394592, + "evolution_kernel_ms": 0.177472, + "evolution_route": "paired_large_v15", + "evolution_speedup": 2.2234, + "evolution_tflops": 48.4016 + } + }, + { + "label": "adjacent_5600_d128_fallback_neighbor_b8_n8320_k256_d128", + "params": { + "B": 8, + "D": 128, + "K": 256, + "N": 8320, + "dtype": "bfloat16", + "seed": 5601281, + "source": "guard_miss_fallback" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 14.6812, + "evolution_flashlib_ms": 0.297119, + "evolution_kernel_ms": 0.155872, + "evolution_route": "aligned_v10_fallback", + "evolution_speedup": 1.9062, + "evolution_tflops": 27.985 + } + }, + { + "label": "post_d895_d144_b1_n256_k256_d144", + "params": { + "B": 1, + "D": 144, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "seed": 214401, + "source": "tailpad_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0351, + "evolution_flashlib_ms": 0.537279, + "evolution_kernel_ms": 0.167728, + "evolution_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_speedup": 3.2033, + "evolution_tflops": 0.1125 + } + }, + { + "label": "post_d895_d144_b1_n512_k8192_d144", + "params": { + "B": 1, + "D": 144, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 214404, + "source": "high_k_low_n" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.6035, + "evolution_flashlib_ms": 2.001693, + "evolution_kernel_ms": 0.302528, + "evolution_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_speedup": 6.6166, + "evolution_tflops": 3.9929 + } + }, + { + "label": "post_d895_d144_b2_n2048_k1024_d144", + "params": { + "B": 2, + "D": 144, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "seed": 214402, + "source": "tailpad_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.5941, + "evolution_flashlib_ms": 0.465663, + "evolution_kernel_ms": 0.192896, + "evolution_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_speedup": 2.4141, + "evolution_tflops": 6.2622 + } + }, + { + "label": "post_d895_d144_b4_n8192_k1024_d144", + "params": { + "B": 4, + "D": 144, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "seed": 214403, + "source": "tailpad_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 20.3965, + "evolution_flashlib_ms": 0.473791, + "evolution_kernel_ms": 0.195104, + "evolution_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_speedup": 2.4284, + "evolution_tflops": 49.5309 + } + }, + { + "label": "post_d895_d176_b1_n256_k256_d176", + "params": { + "B": 1, + "D": 176, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "seed": 217601, + "source": "tailpad_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0453, + "evolution_flashlib_ms": 0.509327, + "evolution_kernel_ms": 0.169792, + "evolution_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_speedup": 2.9997, + "evolution_tflops": 0.1359 + } + }, + { + "label": "post_d895_d176_b1_n512_k8192_d176", + "params": { + "B": 1, + "D": 176, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 217604, + "source": "high_k_low_n" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.72, + "evolution_flashlib_ms": 2.050605, + "evolution_kernel_ms": 0.301759, + "evolution_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_speedup": 6.7955, + "evolution_tflops": 4.8926 + } + }, + { + "label": "post_d895_d176_b2_n2048_k1024_d176", + "params": { + "B": 2, + "D": 176, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "seed": 217602, + "source": "tailpad_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 3.0356, + "evolution_flashlib_ms": 0.486367, + "evolution_kernel_ms": 0.191839, + "evolution_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_speedup": 2.5353, + "evolution_tflops": 7.696 + } + }, + { + "label": "post_d895_d176_b4_n8192_k1024_d176", + "params": { + "B": 4, + "D": 176, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "seed": 217603, + "source": "tailpad_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 23.9979, + "evolution_flashlib_ms": 0.492175, + "evolution_kernel_ms": 0.195712, + "evolution_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "evolution_speedup": 2.5148, + "evolution_tflops": 60.3497 + } + }, + { + "label": "post_d895_d224_b1_n256_k256_d224", + "params": { + "B": 1, + "D": 224, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "seed": 222404, + "source": "medium_small_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0611, + "evolution_flashlib_ms": 0.480639, + "evolution_kernel_ms": 0.170496, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.8191, + "evolution_tflops": 0.1722 + } + }, + { + "label": "post_d895_d224_b1_n512_k8192_d224", + "params": { + "B": 1, + "D": 224, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 222403, + "source": "high_k_low_n" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.9081, + "evolution_flashlib_ms": 2.069245, + "evolution_kernel_ms": 0.236528, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 8.7484, + "evolution_tflops": 7.9443 + } + }, + { + "label": "post_d895_d224_b2_n2048_k1024_d224", + "params": { + "B": 2, + "D": 224, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "seed": 222401, + "source": "new_medium_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 3.9885, + "evolution_flashlib_ms": 0.47112, + "evolution_kernel_ms": 0.176704, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.6662, + "evolution_tflops": 10.6339 + } + }, + { + "label": "adjacent_d9d5_d224_random_midk_b2_n2944_k1280_d224", + "params": { + "B": 2, + "D": 224, + "K": 1280, + "N": 2944, + "dtype": "bfloat16", + "seed": 9522401, + "source": "random_legal" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 6.9153, + "evolution_flashlib_ms": 0.488255, + "evolution_kernel_ms": 0.180384, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.7068, + "evolution_tflops": 18.7179 + } + }, + { + "label": "adjacent_9ca0_d224_guard_overlap_b2_n4096_k256_d224", + "params": { + "B": 2, + "D": 224, + "K": 256, + "N": 4096, + "dtype": "bfloat16", + "seed": 922405, + "source": "guard_overlap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 3.6395, + "evolution_flashlib_ms": 0.258143, + "evolution_kernel_ms": 0.1712, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.5078, + "evolution_tflops": 5.4879 + } + }, + { + "label": "adjacent_ef00_d224_midn_overlap_b2_n5376_k512_d224", + "params": { + "B": 2, + "D": 224, + "K": 512, + "N": 5376, + "dtype": "bfloat16", + "seed": 9002241, + "source": "guard_overlap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 6.6916, + "evolution_flashlib_ms": 0.36856, + "evolution_kernel_ms": 0.177808, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.0728, + "evolution_tflops": 13.8703 + } + }, + { + "label": "adjacent_a2f8_d224_highn_overlap_b2_n6144_k768_d224", + "params": { + "B": 2, + "D": 224, + "K": 768, + "N": 6144, + "dtype": "bfloat16", + "seed": 1028224, + "source": "guard_overlap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 9.3487, + "evolution_flashlib_ms": 0.45224, + "evolution_kernel_ms": 0.176256, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.5658, + "evolution_tflops": 23.987 + } + }, + { + "label": "adjacent_5600_d224_k512_overlap_b3_n3072_k512_d224", + "params": { + "B": 3, + "D": 224, + "K": 512, + "N": 3072, + "dtype": "bfloat16", + "seed": 5602241, + "source": "guard_overlap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 6.148, + "evolution_flashlib_ms": 0.343839, + "evolution_kernel_ms": 0.174928, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.9656, + "evolution_tflops": 12.0846 + } + }, + { + "label": "adjacent_3328_d224_tail_div_b3_n3840_k512_d224", + "params": { + "B": 3, + "D": 224, + "K": 512, + "N": 3840, + "dtype": "bfloat16", + "seed": 3328224, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 8.3041, + "evolution_flashlib_ms": 0.318207, + "evolution_kernel_ms": 0.174895, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.8194, + "evolution_tflops": 15.1086 + } + }, + { + "label": "adjacent_68cf_d224_overlap_b3_n5120_k1024_d224", + "params": { + "B": 3, + "D": 224, + "K": 1024, + "N": 5120, + "dtype": "bfloat16", + "seed": 6822401, + "source": "guard_overlap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 14.4926, + "evolution_flashlib_ms": 0.486208, + "evolution_kernel_ms": 0.180032, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.7007, + "evolution_tflops": 39.1399 + } + }, + { + "label": "adjacent_8f09_d224_low_n_boundary_b4_n1536_k256_d224", + "params": { + "B": 4, + "D": 224, + "K": 256, + "N": 1536, + "dtype": "bfloat16", + "seed": 8092241, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.2208, + "evolution_flashlib_ms": 0.317296, + "evolution_kernel_ms": 0.173664, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.8271, + "evolution_tflops": 4.0575 + } + }, + { + "label": "adjacent_c44f_d224_overlap_b4_n4480_k512_d224", + "params": { + "B": 4, + "D": 224, + "K": 512, + "N": 4480, + "dtype": "bfloat16", + "seed": 4422401, + "source": "guard_overlap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 10.201, + "evolution_flashlib_ms": 0.402944, + "evolution_kernel_ms": 0.173008, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.329, + "evolution_tflops": 23.7585 + } + }, + { + "label": "post_d895_d224_b4_n8192_k1024_d224", + "params": { + "B": 4, + "D": 224, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "seed": 222402, + "source": "new_medium_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 27.1681, + "evolution_flashlib_ms": 0.553311, + "evolution_kernel_ms": 0.190783, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.9002, + "evolution_tflops": 78.7929 + } + }, + { + "label": "adjacent_1d49_d224_highn_overlap_b5_n5632_k768_d224", + "params": { + "B": 5, + "D": 224, + "K": 768, + "N": 5632, + "dtype": "bfloat16", + "seed": 1492241, + "source": "guard_overlap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 20.8173, + "evolution_flashlib_ms": 0.465423, + "evolution_kernel_ms": 0.186336, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.4978, + "evolution_tflops": 51.9966 + } + }, + { + "label": "post_d895_d288_b1_n256_k256_d288", + "params": { + "B": 1, + "D": 288, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "seed": 228804, + "source": "medium_small_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0649, + "evolution_flashlib_ms": 0.581567, + "evolution_kernel_ms": 0.171104, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.3989, + "evolution_tflops": 0.2206 + } + }, + { + "label": "adjacent_9ca0_d288_highk_low_n_b1_n384_k4096_d288", + "params": { + "B": 1, + "D": 288, + "K": 4096, + "N": 384, + "dtype": "bfloat16", + "seed": 928805, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0159, + "evolution_flashlib_ms": 56.807529, + "evolution_kernel_ms": 0.243872, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 232.9399, + "evolution_tflops": 3.7149 + } + }, + { + "label": "post_d895_d288_b1_n512_k8192_d288", + "params": { + "B": 1, + "D": 288, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 228803, + "source": "high_k_low_n" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0421, + "evolution_flashlib_ms": 57.444236, + "evolution_kernel_ms": 0.273615, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 209.9455, + "evolution_tflops": 8.8296 + } + }, + { + "label": "adjacent_a2f8_d288_splitk_probe_b1_n640_k4096_d288", + "params": { + "B": 1, + "D": 288, + "K": 4096, + "N": 640, + "dtype": "bfloat16", + "seed": 1028288, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0262, + "evolution_flashlib_ms": 57.6586, + "evolution_kernel_ms": 0.24416, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 236.1509, + "evolution_tflops": 6.1843 + } + }, + { + "label": "adjacent_d9d5_d288_heldout_low_n_b1_n896_k4096_d288", + "params": { + "B": 1, + "D": 288, + "K": 4096, + "N": 896, + "dtype": "bfloat16", + "seed": 9528801, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0365, + "evolution_flashlib_ms": 57.939113, + "evolution_kernel_ms": 0.2448, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 236.6794, + "evolution_tflops": 8.6353 + } + }, + { + "label": "adjacent_ef00_d288_highk_heldout_b1_n1664_k4096_d288", + "params": { + "B": 1, + "D": 288, + "K": 4096, + "N": 1664, + "dtype": "bfloat16", + "seed": 9002881, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0667, + "evolution_flashlib_ms": 58.88955, + "evolution_kernel_ms": 0.247408, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 238.0261, + "evolution_tflops": 15.868 + } + }, + { + "label": "adjacent_8f09_d288_k768_overlap_b1_n2560_k768_d288", + "params": { + "B": 1, + "D": 288, + "K": 768, + "N": 2560, + "dtype": "bfloat16", + "seed": 8092881, + "source": "guard_overlap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.3579, + "evolution_flashlib_ms": 0.480287, + "evolution_kernel_ms": 0.18432, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.6057, + "evolution_tflops": 6.144 + } + }, + { + "label": "adjacent_c44f_d288_highk_heldout_b2_n768_k4096_d288", + "params": { + "B": 2, + "D": 288, + "K": 4096, + "N": 768, + "dtype": "bfloat16", + "seed": 4428801, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.2892, + "evolution_flashlib_ms": 1.583007, + "evolution_kernel_ms": 0.246496, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 6.422, + "evolution_tflops": 14.7016 + } + }, + { + "label": "adjacent_5600_d288_mid_highk_b2_n1024_k2048_d288", + "params": { + "B": 2, + "D": 288, + "K": 2048, + "N": 1024, + "dtype": "bfloat16", + "seed": 5602881, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.4751, + "evolution_flashlib_ms": 0.976096, + "evolution_kernel_ms": 0.206512, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 4.7266, + "evolution_tflops": 11.6987 + } + }, + { + "label": "adjacent_68cf_d288_boundary_b2_n1920_k512_d288", + "params": { + "B": 2, + "D": 288, + "K": 512, + "N": 1920, + "dtype": "bfloat16", + "seed": 6828801, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.6899, + "evolution_flashlib_ms": 0.421007, + "evolution_kernel_ms": 0.177824, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.3676, + "evolution_tflops": 6.3684 + } + }, + { + "label": "post_d895_d288_b2_n2048_k1024_d288", + "params": { + "B": 2, + "D": 288, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "seed": 228801, + "source": "new_medium_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 3.7159, + "evolution_flashlib_ms": 0.65016, + "evolution_kernel_ms": 0.18616, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.4925, + "evolution_tflops": 12.9777 + } + }, + { + "label": "adjacent_1d49_d288_low_n_heldout_b3_n1152_k2048_d288", + "params": { + "B": 3, + "D": 288, + "K": 2048, + "N": 1152, + "dtype": "bfloat16", + "seed": 1492881, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 4.0485, + "evolution_flashlib_ms": 1.007006, + "evolution_kernel_ms": 0.207744, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 4.8473, + "evolution_tflops": 19.6245 + } + }, + { + "label": "adjacent_3328_d288_guard_overlap_b4_n8192_k256_d288", + "params": { + "B": 4, + "D": 288, + "K": 256, + "N": 8192, + "dtype": "bfloat16", + "seed": 3328288, + "source": "guard_overlap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 14.007, + "evolution_flashlib_ms": 0.344959, + "evolution_kernel_ms": 0.182432, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.8909, + "evolution_tflops": 26.4857 + } + }, + { + "label": "post_d895_d288_b4_n8192_k1024_d288", + "params": { + "B": 4, + "D": 288, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "seed": 228802, + "source": "new_medium_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 28.8551, + "evolution_flashlib_ms": 0.669808, + "evolution_kernel_ms": 0.214687, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.1199, + "evolution_tflops": 90.0257 + } + }, + { + "label": "post_d895_d352_b1_n256_k256_d352", + "params": { + "B": 1, + "D": 352, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "seed": 235204, + "source": "high_small_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0736, + "evolution_flashlib_ms": 0.627167, + "evolution_kernel_ms": 0.172703, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.6315, + "evolution_tflops": 0.2671 + } + }, + { + "label": "adjacent_9ca0_d352_splitk_boundary_b1_n512_k8192_d352", + "params": { + "B": 1, + "D": 352, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 935205, + "source": "request_specific" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0601, + "evolution_flashlib_ms": 49.145822, + "evolution_kernel_ms": 0.275519, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 178.3754, + "evolution_tflops": 10.7172 + } + }, + { + "label": "post_d895_d352_b1_n512_k8192_d352", + "params": { + "B": 1, + "D": 352, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 235203, + "source": "high_k_low_n" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0616, + "evolution_flashlib_ms": 47.910581, + "evolution_kernel_ms": 0.272767, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 175.6465, + "evolution_tflops": 10.8253 + } + }, + { + "label": "adjacent_c44f_d352_random_b1_n3328_k768_d352", + "params": { + "B": 1, + "D": 352, + "K": 768, + "N": 3328, + "dtype": "bfloat16", + "seed": 4435201, + "source": "random_legal" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 3.6993, + "evolution_flashlib_ms": 0.4864, + "evolution_kernel_ms": 0.187632, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.5923, + "evolution_tflops": 9.5898 + } + }, + { + "label": "adjacent_5600_d352_splitk_edge_b2_n768_k4096_d352", + "params": { + "B": 2, + "D": 352, + "K": 4096, + "N": 768, + "dtype": "bfloat16", + "seed": 5603521, + "source": "request_specific" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.7323, + "evolution_flashlib_ms": 1.621022, + "evolution_kernel_ms": 1.051615, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.5415, + "evolution_tflops": 4.2118 + } + }, + { + "label": "adjacent_1d49_d352_request_highk_b2_n1024_k8192_d352", + "params": { + "B": 2, + "D": 352, + "K": 8192, + "N": 1024, + "dtype": "bfloat16", + "seed": 1493521, + "source": "request_specific" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 4.5385, + "evolution_flashlib_ms": 2.602461, + "evolution_kernel_ms": 0.302655, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 8.5988, + "evolution_tflops": 39.0252 + } + }, + { + "label": "post_d895_d352_b2_n2048_k1024_d352", + "params": { + "B": 2, + "D": 352, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "seed": 235201, + "source": "new_high_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 4.882, + "evolution_flashlib_ms": 0.604831, + "evolution_kernel_ms": 0.189184, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.1971, + "evolution_tflops": 15.608 + } + }, + { + "label": "adjacent_ef00_d352_request_low_n_b3_n896_k8192_d352", + "params": { + "B": 3, + "D": 352, + "K": 8192, + "N": 896, + "dtype": "bfloat16", + "seed": 9003521, + "source": "request_specific" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 4.9423, + "evolution_flashlib_ms": 3.136652, + "evolution_kernel_ms": 0.316191, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 9.9201, + "evolution_tflops": 49.0278 + } + }, + { + "label": "adjacent_3328_d352_random_legal_b3_n2048_k768_d352", + "params": { + "B": 3, + "D": 352, + "K": 768, + "N": 2048, + "dtype": "bfloat16", + "seed": 3328352, + "source": "random_legal" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 5.7044, + "evolution_flashlib_ms": 0.582335, + "evolution_kernel_ms": 0.183872, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.1671, + "evolution_tflops": 18.0663 + } + }, + { + "label": "adjacent_68cf_d352_tail_b3_n2816_k768_d352", + "params": { + "B": 3, + "D": 352, + "K": 768, + "N": 2816, + "dtype": "bfloat16", + "seed": 6835201, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 9.1314, + "evolution_flashlib_ms": 0.500207, + "evolution_kernel_ms": 0.18768, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.6652, + "evolution_tflops": 24.3372 + } + }, + { + "label": "adjacent_a2f8_d352_smallk_boundary_b4_n1024_k256_d352", + "params": { + "B": 4, + "D": 352, + "K": 256, + "N": 1024, + "dtype": "bfloat16", + "seed": 1028352, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 1.5395, + "evolution_flashlib_ms": 0.479503, + "evolution_kernel_ms": 0.176096, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.723, + "evolution_tflops": 4.192 + } + }, + { + "label": "adjacent_8f09_d352_smallk_large_n_b4_n4096_k256_d352", + "params": { + "B": 4, + "D": 352, + "K": 256, + "N": 4096, + "dtype": "bfloat16", + "seed": 8093521, + "source": "random_legal" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 9.629, + "evolution_flashlib_ms": 0.306655, + "evolution_kernel_ms": 0.176128, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.7411, + "evolution_tflops": 16.765 + } + }, + { + "label": "post_d895_d352_b4_n8192_k1024_d352", + "params": { + "B": 4, + "D": 352, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "seed": 235202, + "source": "new_high_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 36.2395, + "evolution_flashlib_ms": 0.651839, + "evolution_kernel_ms": 0.220992, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.9496, + "evolution_tflops": 106.8922 + } + }, + { + "label": "adjacent_d9d5_d352_random_b5_n2304_k768_d352", + "params": { + "B": 5, + "D": 352, + "K": 768, + "N": 2304, + "dtype": "bfloat16", + "seed": 9535201, + "source": "random_legal" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 10.6911, + "evolution_flashlib_ms": 0.582592, + "evolution_kernel_ms": 0.186944, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.1164, + "evolution_tflops": 33.3177 + } + }, + { + "label": "post_d895_d416_b1_n256_k256_d416", + "params": { + "B": 1, + "D": 416, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "seed": 241604, + "source": "high_small_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0842, + "evolution_flashlib_ms": 0.647807, + "evolution_kernel_ms": 0.175137, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.6989, + "evolution_tflops": 0.3113 + } + }, + { + "label": "adjacent_3328_d416_highk_low_n_b1_n384_k8192_d416", + "params": { + "B": 1, + "D": 416, + "K": 8192, + "N": 384, + "dtype": "bfloat16", + "seed": 3328416, + "source": "request_specific" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.063, + "evolution_flashlib_ms": 41.538691, + "evolution_kernel_ms": 0.276031, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 150.4856, + "evolution_tflops": 9.4817 + } + }, + { + "label": "post_d895_d416_b1_n512_k8192_d416", + "params": { + "B": 1, + "D": 416, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 241603, + "source": "high_k_low_n" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0852, + "evolution_flashlib_ms": 40.967689, + "evolution_kernel_ms": 0.275103, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 148.9176, + "evolution_tflops": 12.6849 + } + }, + { + "label": "adjacent_9ca0_d416_splitk_min_b1_n1024_k4096_d416", + "params": { + "B": 1, + "D": 416, + "K": 4096, + "N": 1024, + "dtype": "bfloat16", + "seed": 941605, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0848, + "evolution_flashlib_ms": 41.154087, + "evolution_kernel_ms": 0.274271, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 150.049, + "evolution_tflops": 12.7234 + } + }, + { + "label": "adjacent_ef00_d416_random_midk_b1_n2176_k768_d416", + "params": { + "B": 1, + "D": 416, + "K": 768, + "N": 2176, + "dtype": "bfloat16", + "seed": 9004161, + "source": "random_legal" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.6268, + "evolution_flashlib_ms": 0.529311, + "evolution_kernel_ms": 0.191648, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.7619, + "evolution_tflops": 7.255 + } + }, + { + "label": "adjacent_d9d5_d416_highk_request_b2_n640_k8192_d416", + "params": { + "B": 2, + "D": 416, + "K": 8192, + "N": 640, + "dtype": "bfloat16", + "seed": 9541601, + "source": "request_specific" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.6807, + "evolution_flashlib_ms": 3.254428, + "evolution_kernel_ms": 0.298432, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 10.9051, + "evolution_tflops": 29.2333 + } + }, + { + "label": "post_d895_d416_b2_n2048_k1024_d416", + "params": { + "B": 2, + "D": 416, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "seed": 241601, + "source": "new_high_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 5.4317, + "evolution_flashlib_ms": 0.642463, + "evolution_kernel_ms": 0.193983, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.312, + "evolution_tflops": 17.9895 + } + }, + { + "label": "adjacent_68cf_d416_overlap_b2_n2304_k1024_d416", + "params": { + "B": 2, + "D": 416, + "K": 1024, + "N": 2304, + "dtype": "bfloat16", + "seed": 6841601, + "source": "guard_overlap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 6.4021, + "evolution_flashlib_ms": 0.613215, + "evolution_kernel_ms": 0.196895, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.1144, + "evolution_tflops": 19.9389 + } + }, + { + "label": "adjacent_a2f8_d416_random_b2_n2560_k768_d416", + "params": { + "B": 2, + "D": 416, + "K": 768, + "N": 2560, + "dtype": "bfloat16", + "seed": 1028416, + "source": "random_legal" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 6.6441, + "evolution_flashlib_ms": 0.492399, + "evolution_kernel_ms": 0.189631, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.5966, + "evolution_tflops": 17.2522 + } + }, + { + "label": "adjacent_c44f_d416_highk_request_b3_n512_k8192_d416", + "params": { + "B": 3, + "D": 416, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 4441601, + "source": "request_specific" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.7386, + "evolution_flashlib_ms": 3.822732, + "evolution_kernel_ms": 0.303423, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 12.5987, + "evolution_tflops": 34.5029 + } + }, + { + "label": "adjacent_8f09_d416_midk_tail_b3_n3456_k768_d416", + "params": { + "B": 3, + "D": 416, + "K": 768, + "N": 3456, + "dtype": "bfloat16", + "seed": 8094161, + "source": "tail_divisibility" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 13.4164, + "evolution_flashlib_ms": 0.493791, + "evolution_kernel_ms": 0.189088, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.6114, + "evolution_tflops": 35.0361 + } + }, + { + "label": "adjacent_5600_d416_smallk_boundary_b4_n2048_k256_d416", + "params": { + "B": 4, + "D": 416, + "K": 256, + "N": 2048, + "dtype": "bfloat16", + "seed": 5604161, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 4.7141, + "evolution_flashlib_ms": 0.370128, + "evolution_kernel_ms": 0.175743, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.1061, + "evolution_tflops": 9.9283 + } + }, + { + "label": "adjacent_1d49_d416_random_b4_n3840_k512_d416", + "params": { + "B": 4, + "D": 416, + "K": 512, + "N": 3840, + "dtype": "bfloat16", + "seed": 1494161, + "source": "random_legal" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 16.9546, + "evolution_flashlib_ms": 0.385919, + "evolution_kernel_ms": 0.184, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.0974, + "evolution_tflops": 35.5604 + } + }, + { + "label": "post_d895_d416_b4_n8192_k1024_d416", + "params": { + "B": 4, + "D": 416, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "seed": 241602, + "source": "new_high_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 39.3193, + "evolution_flashlib_ms": 0.710015, + "evolution_kernel_ms": 0.230463, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.0808, + "evolution_tflops": 121.1354 + } + }, + { + "label": "adjacent_3328_d480_min_boundary_b1_n128_k256_d480", + "params": { + "B": 1, + "D": 480, + "K": 256, + "N": 128, + "dtype": "bfloat16", + "seed": 3328480, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0644, + "evolution_flashlib_ms": 0.488575, + "evolution_kernel_ms": 0.177056, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.7594, + "evolution_tflops": 0.1777 + } + }, + { + "label": "post_d895_d480_b1_n256_k256_d480", + "params": { + "B": 1, + "D": 480, + "K": 256, + "N": 256, + "dtype": "bfloat16", + "seed": 248004, + "source": "high_small_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.0969, + "evolution_flashlib_ms": 0.649151, + "evolution_kernel_ms": 0.176655, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.6747, + "evolution_tflops": 0.3561 + } + }, + { + "label": "post_d895_d480_b1_n512_k8192_d480", + "params": { + "B": 1, + "D": 480, + "K": 8192, + "N": 512, + "dtype": "bfloat16", + "seed": 248003, + "source": "high_k_low_n" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.1226, + "evolution_flashlib_ms": 32.836689, + "evolution_kernel_ms": 0.278559, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 117.8803, + "evolution_tflops": 14.4548 + } + }, + { + "label": "adjacent_a2f8_d480_highk_request_b1_n896_k4096_d480", + "params": { + "B": 1, + "D": 480, + "K": 4096, + "N": 896, + "dtype": "bfloat16", + "seed": 1028480, + "source": "request_specific" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 0.1098, + "evolution_flashlib_ms": 32.096809, + "evolution_kernel_ms": 0.274432, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 116.9572, + "evolution_tflops": 12.8382 + } + }, + { + "label": "adjacent_5600_d480_random_b1_n1536_k1024_d480", + "params": { + "B": 1, + "D": 480, + "K": 1024, + "N": 1536, + "dtype": "bfloat16", + "seed": 5604801, + "source": "random_legal" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.1001, + "evolution_flashlib_ms": 0.718975, + "evolution_kernel_ms": 0.199488, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.6041, + "evolution_tflops": 7.5691 + } + }, + { + "label": "adjacent_8f09_d480_highk_low_n_b2_n640_k4096_d480", + "params": { + "B": 2, + "D": 480, + "K": 4096, + "N": 640, + "dtype": "bfloat16", + "seed": 8094801, + "source": "request_specific" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 2.5174, + "evolution_flashlib_ms": 1.999325, + "evolution_kernel_ms": 0.284192, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 7.0351, + "evolution_tflops": 17.7104 + } + }, + { + "label": "post_d895_d480_b2_n2048_k1024_d480", + "params": { + "B": 2, + "D": 480, + "K": 1024, + "N": 2048, + "dtype": "bfloat16", + "seed": 248001, + "source": "new_high_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 5.9879, + "evolution_flashlib_ms": 0.672447, + "evolution_kernel_ms": 0.197919, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.3976, + "evolution_tflops": 20.3443 + } + }, + { + "label": "adjacent_1d49_d480_smallk_boundary_b2_n2816_k256_d480", + "params": { + "B": 2, + "D": 480, + "K": 256, + "N": 2816, + "dtype": "bfloat16", + "seed": 1494801, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 4.1626, + "evolution_flashlib_ms": 0.332511, + "evolution_kernel_ms": 0.178015, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.8679, + "evolution_tflops": 7.7753 + } + }, + { + "label": "adjacent_9ca0_d480_splitd_large_n_b2_n4096_k512_d480", + "params": { + "B": 2, + "D": 480, + "K": 512, + "N": 4096, + "dtype": "bfloat16", + "seed": 948005, + "source": "random_legal" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 8.9917, + "evolution_flashlib_ms": 0.447807, + "evolution_kernel_ms": 0.185087, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.4194, + "evolution_tflops": 21.7548 + } + }, + { + "label": "adjacent_d9d5_d480_heldout_highk_b3_n1024_k4096_d480", + "params": { + "B": 3, + "D": 480, + "K": 4096, + "N": 1024, + "dtype": "bfloat16", + "seed": 9548001, + "source": "heldout_neighborhood" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 6.283, + "evolution_flashlib_ms": 1.92259, + "evolution_kernel_ms": 0.297632, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 6.4596, + "evolution_tflops": 40.5857 + } + }, + { + "label": "adjacent_ef00_d480_lowk_boundary_b3_n3200_k256_d480", + "params": { + "B": 3, + "D": 480, + "K": 256, + "N": 3200, + "dtype": "bfloat16", + "seed": 9004801, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 7.059, + "evolution_flashlib_ms": 0.334224, + "evolution_kernel_ms": 0.18192, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 1.8372, + "evolution_tflops": 12.9689 + } + }, + { + "label": "adjacent_68cf_d480_boundary_b4_n1664_k512_d480", + "params": { + "B": 4, + "D": 480, + "K": 512, + "N": 1664, + "dtype": "bfloat16", + "seed": 6848001, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 6.0602, + "evolution_flashlib_ms": 0.539839, + "evolution_kernel_ms": 0.18536, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.9124, + "evolution_tflops": 17.6497 + } + }, + { + "label": "post_d895_d480_b4_n8192_k1024_d480", + "params": { + "B": 4, + "D": 480, + "K": 1024, + "N": 8192, + "dtype": "bfloat16", + "seed": 248002, + "source": "new_high_gap" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 43.8869, + "evolution_flashlib_ms": 0.733983, + "evolution_kernel_ms": 0.24, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 3.0583, + "evolution_tflops": 134.2177 + } + }, + { + "label": "adjacent_c44f_d480_boundary_b5_n2048_k512_d480", + "params": { + "B": 5, + "D": 480, + "K": 512, + "N": 2048, + "dtype": "bfloat16", + "seed": 4448001, + "source": "guard_boundary" + }, + "recorded": { + "evolution_flashlib_equiv_tflops": 10.031, + "evolution_flashlib_ms": 0.501759, + "evolution_kernel_ms": 0.185136, + "evolution_route": "gap_pad_to_supported_seed_v1", + "evolution_speedup": 2.7102, + "evolution_tflops": 27.1863 + } + } +] diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_launch_plan.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_launch_plan.py new file mode 100644 index 00000000..8a99cc8c --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_launch_plan.py @@ -0,0 +1,540 @@ +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field +from functools import lru_cache +from typing import Any + +from ._dispatch import flash_kmeans_assign_dispatcher as _root +from ._dispatch_runtime import _import_dispatch_module, dispatch_launch_options + +_WEAVE_PREFIX = 'loom.examples.weave.' +_ROOT_MODULE = 'flash_kmeans_assign_dispatcher' +_ROOT_CALLABLE = 'launch_for_eval' +_EXACT_LAUNCH_SPECS = {} + + +@dataclass(frozen=True) +class RouteDecision: + """Resolved semantic route with a launcher that can be reused directly.""" + + route_id: str + launch_entrypoint: str + launcher: Callable[[dict[str, Any]], Any] = field(repr=False, compare=False) + exact_contract: bool = False + + def launch( + self, + inputs: dict[str, Any], + *, + stream: Any = None, + timeout_ms: float | None = None, + ) -> Any: + with dispatch_launch_options(stream=stream, timeout_ms=timeout_ms): + return self.launcher(inputs) + + +def _route_key(inputs: dict[str, Any]) -> tuple[int, int, int, int, int, str, bool, bool]: + dtype = str(inputs.get("dtype", "bfloat16")) + if dtype.startswith("torch."): + dtype = dtype[6:] + return ( + *(int(inputs[name]) for name in ("B", "Q", "M", "D", "K")), + dtype, + bool(inputs.get("self_search", False)), + bool(inputs.get("force_fallback", False)), + ) + + +@lru_cache(maxsize=None) +def _load_launcher(module_name: str, callable_name: str) -> Callable[[dict[str, Any]], Any]: + module = _import_dispatch_module(module_name) + launcher = getattr(module, callable_name, None) + if not callable(launcher): + raise RuntimeError(f"resolved dispatcher launcher is not callable: {module_name}:{callable_name}") + return launcher + + +@lru_cache(maxsize=None) +def _make_decision( + route_id: str, + module_name: str, + callable_name: str, + exact_contract: bool, +) -> RouteDecision: + return RouteDecision( + route_id=route_id, + launch_entrypoint=f"{_WEAVE_PREFIX}{module_name}:{callable_name}", + launcher=_load_launcher(module_name, callable_name), + exact_contract=exact_contract, + ) + + +def _entrypoint_spec(entrypoint: object) -> tuple[str, str] | None: + if not isinstance(entrypoint, str): + return None + module_name, separator, callable_name = entrypoint.partition(":") + if not separator or not module_name.startswith(_WEAVE_PREFIX) or not callable_name.isidentifier(): + return None + return module_name.removeprefix(_WEAVE_PREFIX), callable_name + + +def _generic_decision(inputs: dict[str, Any]) -> RouteDecision: + info_fn = getattr(_root, "route_info", None) + info = dict(info_fn(inputs)) if callable(info_fn) else {} + route_id = info.get("selected_route", info.get("route")) + if route_id is None: + select_route = getattr(_root, "selected_route", None) + route_id = select_route(inputs) if callable(select_route) else _ROOT_CALLABLE + # ``resolved_launch_entrypoint`` is an explicit launch contract and may + # bypass the root. ``selected_entrypoint`` is seed provenance: it can name + # a narrower module whose own guards would silently miss for this + # signature (the K11 prefix route reports its K64 seed there, and + # launching the seed with K=11 inputs falls through to a slower exact + # parent). Signatures without a launch contract go through the root + # dispatcher, which reproduces the frozen guard cascade exactly. + spec = _entrypoint_spec(info.get("resolved_launch_entrypoint")) + if spec is None: + spec = (_ROOT_MODULE, _ROOT_CALLABLE) + return _make_decision(str(route_id), *spec, False) + + +def resolve_route(inputs: dict[str, Any]) -> RouteDecision: + """Resolve once; exact exported shapes never re-enter the root dispatcher.""" + + spec = _EXACT_LAUNCH_SPECS.get(_route_key(inputs)) + if spec is None: + return _generic_decision(inputs) + return _make_decision(*spec, True) + + +class LaunchPlan: + """Per-signature resolved execution state for the exported hot path. + + Migration step 2 of ``PLAN_BASED_EXPORT_RUNTIME.md``: the guard cascade + (``resolve_route`` plus one captured dispatcher traversal) runs exactly + once, at construction. A hot call overwrites the recorded 8-byte pointer + carriers in place, refreshes device tensor-map descriptors only when a + bound pointer actually changed, and submits the already-marshalled + launches on their construction-time stream — no re-marshalling, no + per-launch stream query, no dispatcher re-entry. + """ + + __slots__ = ( + "route", + "sequence", + "torch_stream", + "stream_handle", + "_launches", + "_pointer_writers", + "_tma_bindings", + ) + + def __init__(self, route: RouteDecision, sequence: Any, *, torch_stream: Any, stream_handle: int): + launches = tuple(sequence._launches) + input_bindings = tuple(sequence._input_bindings) + if input_bindings and len(input_bindings) != len(launches): + raise RuntimeError("launch plan capture has corrupt input bindings") + writers: dict[str, list[Any]] = {} + for launch, bindings in zip(launches, input_bindings): + carriers = launch._packed._prevent_gc + for index, key in bindings: + writers.setdefault(key, []).append(carriers[index]) + self.route = route + self.sequence = sequence + self.torch_stream = torch_stream + self.stream_handle = int(stream_handle) + self._launches = launches + self._pointer_writers = tuple((key, tuple(items)) for key, items in writers.items()) + self._tma_bindings = tuple(sequence._tensor_map_bindings) + + @property + def launch_count(self) -> int: + return len(self._launches) + + @property + def bound_input_keys(self) -> tuple[str, ...]: + direct = {key for key, _carriers in self._pointer_writers} + derived = {binding.input_key for binding in self._tma_bindings} + return tuple(sorted(direct | derived)) + + def bind_hot(self, bindings: dict[str, Any]) -> None: + """Refresh tensor maps and overwrite bound pointer carriers in place. + + This is the host-side half of a hot call. Callers that enqueue their + own support launches between the plan's pointer binding and its + kernel submission (the KNN-build fused row norms) must call this + BEFORE those launches: a fresh-pointer tensor-map re-encode costs + host time, and paying it after any kernel is already enqueued turns + that host time into a GPU inter-kernel gap. + + Tensor maps go through the per-pointer variant bank: a pointer the + bank has seen keeps its device-resident descriptor, so re-activating + it is a handful of carrier writes with no ``cuTensorMapEncodeTiled``, + no pinned-staging refresh, and no H2D copy. The plan's signature slot + is stream-keyed and alias-keyed by the caller, which is the safety + contract ``rebind_stream_bound`` requires. + """ + + for binding in self._tma_bindings: + binding.rebind_stream_bound(bindings[binding.input_key], stream=self.torch_stream) + for key, carriers in self._pointer_writers: + pointer = bindings[key].data_ptr() + for carrier in carriers: + carrier.value = pointer + + def submit_hot(self, *, timeout_ms: float | None = None) -> None: + """Submit every prepared launch on the plan's construction stream.""" + + launches = self._launches + last_index = len(launches) - 1 + for index, launch in enumerate(launches): + launch.launch(stream=None, timeout_ms=timeout_ms if index == last_index else None) + + def launch_hot(self, bindings: dict[str, Any], *, timeout_ms: float | None = None) -> None: + """Patch bound pointer carriers from ``bindings`` and submit every launch.""" + + self.bind_hot(bindings) + self.submit_hot(timeout_ms=timeout_ms) + + def record_stream(self, stream: Any) -> None: + """Tie every plan-held launch argument to ``stream`` before release.""" + + self.sequence.record_stream(stream) + + +class PerCallRoutePlan: + """Per-signature plan for routes whose host logic reads device results. + + Capture observed the route reading GPU memory (for example an + ``overflow_flag.item()`` certification) while its kernels were only being + recorded, so a frozen launch list cannot reproduce the route's per-call + branch decisions — freezing would bake whichever branch the capture-time + garbage selected. These signatures keep the generic per-call launcher: + route resolution stays cached and the guard cascade still ran exactly + once, but every hot call re-executes the resolved route's host program. + Device-side repair (Migration step 3) makes such routes replayable. + """ + + __slots__ = ( + "route", + "torch_stream", + "stream_handle", + "launch_count", + "host_data_reads", + "_static_inputs", + "_pending_bindings", + ) + + def __init__( + self, + route: RouteDecision, + *, + torch_stream: Any, + stream_handle: int, + static_inputs: dict[str, Any], + launch_count: int, + host_data_reads: int, + ): + self.route = route + self.torch_stream = torch_stream + self.stream_handle = int(stream_handle) + self.launch_count = int(launch_count) + self.host_data_reads = int(host_data_reads) + self._static_inputs = dict(static_inputs) + self._pending_bindings = None + + def bind_hot(self, bindings: dict[str, Any]) -> None: + """Stage this call's tensor bindings for ``submit_hot``. + + Mirrors ``LaunchPlan``'s two-phase hot call so callers with support + launches use one code path. The caller's per-signature slot lock + serializes bind/submit pairs on a plan instance. + """ + + self._pending_bindings = dict(bindings) + + def submit_hot(self, *, timeout_ms: float | None = None) -> None: + """Re-execute the resolved route with the staged tensor bindings.""" + + bindings = self._pending_bindings + if bindings is None: + raise RuntimeError("PerCallRoutePlan.submit_hot requires a preceding bind_hot") + self._pending_bindings = None + self.launch_hot(bindings, timeout_ms=timeout_ms) + + def launch_hot(self, bindings: dict[str, Any], *, timeout_ms: float | None = None) -> None: + """Re-execute the resolved route with this call's tensor bindings. + + The torch stream context keeps the route's tensor operations (scratch + fills, the certification read-back) ordered with its kernel launches + on the plan's stream, exactly as the live evaluation path runs it. + """ + + import torch + + inputs = dict(self._static_inputs) + inputs.update(bindings) + with torch.cuda.stream(self.torch_stream): + self.route.launch(inputs, stream=self.torch_stream, timeout_ms=timeout_ms) + + def record_stream(self, stream: Any) -> None: + """Per-call plans hold no launch arguments; nothing to record.""" + + +def build_launch_plan( + inputs: dict[str, Any], + *, + stream: Any, + arch: str, + validate_result: Callable[[Any, dict[str, Any]], None] | None = None, + route: Any = None, +) -> LaunchPlan | PerCallRoutePlan: + """Run the guard cascade once and freeze its launches into a LaunchPlan. + + This is the per-signature slow path; ``resolve_route`` remains the single + source of routing truth. ``validate_result(result, inputs)`` must raise + when the resolved route's outputs cannot be retargeted by pointer + rebinding (for example outputs that are not caller-owned tensors). + Routes whose host logic read device memory during capture resolve to a + ``PerCallRoutePlan`` instead of a frozen launch list. + + ``route`` accepts an already-resolved decision from a sibling routing + layer with the same contract (``route_id``/``launch_entrypoint``/ + ``exact_contract`` plus ``launch(inputs, stream=..., timeout_ms=...)``), + for workloads whose exact-contract table lives outside this module (the + KNN-build direct-manifest resolver). It must come from that workload's + frozen routing surface, never from re-guessing the cascade. + """ + + from ._dispatch_runtime import capture_kernel_launches + from ._runtime import launch_context + + if stream is None: + raise ValueError("build_launch_plan requires a resolved torch CUDA stream, not None") + if route is None: + route = resolve_route(inputs) + with capture_kernel_launches(stream=stream, arch=arch, inputs=inputs) as captured: + with launch_context(arch=arch, stream=stream, timeout_ms=None): + result = route.launch(inputs, stream=stream, timeout_ms=None) + if validate_result is not None: + validate_result(result, inputs) + if captured.host_data_dependent: + static_inputs = { + key: value for key, value in inputs.items() if not callable(getattr(value, "data_ptr", None)) + } + return PerCallRoutePlan( + route, + torch_stream=stream, + stream_handle=int(stream.cuda_stream), + static_inputs=static_inputs, + launch_count=len(captured._launches), + host_data_reads=captured.host_data_reads, + ) + sequence = captured.bind(result) + return LaunchPlan( + route, + sequence, + torch_stream=stream, + stream_handle=int(stream.cuda_stream), + ) + + +class GraphCaptureUnsupported(RuntimeError): + """A plan's launches have no validated CUDA-graph capture path.""" + + +def _check_cu(err: Any, message: str) -> None: + code = err[0] if isinstance(err, tuple) else err + if int(code) != 0: + raise RuntimeError(f"{message}: CUresult={int(code)}") + + +class GraphExecPlan: + """One per-signature CUDA graph over a LaunchPlan plus support launches. + + Migration step 3 of ``PLAN_BASED_EXPORT_RUNTIME.md``: the signature's + stable kernel chain (support launches such as fused row norms, then the + frozen route launches) is stream-captured once at plan construction. A + hot call is host-only binding (the caller's ``plan.bind_hot`` plus + support pointer writes into the same persistent packed argument buffers), + then ``submit_hot``: every kernel node's packed buffer is pushed through + ``cuGraphExecKernelNodeSetParams`` and the chain replays with one + ``cuGraphLaunch`` on the plan's construction-time stream. Kernel-node + launch attributes recorded at capture (cluster dimensions, scheduling + preference) persist across exec-node parameter updates. + """ + + __slots__ = ( + "plan", + "_launches", + "_graph", + "_graph_exec", + "_node_params", + "_cu_stream", + "_set_params", + "_graph_launch", + "_cu_success", + "_destroyed", + ) + + def __init__(self, plan: LaunchPlan, launches, graph, graph_exec, node_params, cu_stream): + from cuda.bindings import driver + + self.plan = plan + self._launches = tuple(launches) + self._graph = graph + self._graph_exec = graph_exec + self._node_params = tuple(node_params) + self._cu_stream = cu_stream + self._set_params = driver.cuGraphExecKernelNodeSetParams + self._graph_launch = driver.cuGraphLaunch + self._cu_success = driver.CUresult.CUDA_SUCCESS + self._destroyed = False + + @property + def launch_count(self) -> int: + return len(self._launches) + + def submit_hot(self, *, timeout_ms: float | None = None) -> None: + """Push the persistent packed argument buffers and replay the graph. + + The caller must have completed every pointer/tensor-map bind for this + call (``plan.bind_hot`` plus support-launch binds) first; parameter + values are copied out of the packed buffers here. + """ + + if self._destroyed: + raise RuntimeError("graph plan was destroyed by a runtime clear()") + set_params = self._set_params + graph_exec = self._graph_exec + success = self._cu_success + for node, params in self._node_params: + (err,) = set_params(graph_exec, node, params) + if err != success: + _check_cu(err, "cuGraphExecKernelNodeSetParams failed") + (err,) = self._graph_launch(graph_exec, self._cu_stream) + if err != success: + _check_cu(err, "cuGraphLaunch failed") + if timeout_ms is not None: + self._launches[-1]._kernel._wait_with_timeout(self._cu_stream, timeout_ms) + + def destroy(self) -> None: + """Release the driver graph handles. Device work must be complete.""" + + if self._destroyed: + return + from cuda.bindings import driver + + self._destroyed = True + driver.cuGraphExecDestroy(self._graph_exec) + driver.cuGraphDestroy(self._graph) + + +def build_graph_exec_plan(plan: Any, *, support_launches: tuple = ()) -> GraphExecPlan: + """Capture ``support_launches`` then ``plan``'s launches into one graph. + + The per-signature slow path, run once at plan construction. Launches are + replayed onto a dedicated capture stream (graph construction only — no + kernel executes), each launch is mapped to its kernel node through the + capture stream's leaf-dependency query, and the captured topology is + hard-checked to contain exactly the expected kernel nodes so foreign + work (for example watchdog event records) can never silently ride along. + + Raises ``GraphCaptureUnsupported`` for plans that cannot replay from a + frozen kernel chain (per-call routes) and for launch modes without a + validated capture path (cooperative). Any other failure propagates — + a capture that should work but does not is an error, not a fallback. + """ + + import ctypes + import sys + from contextlib import nullcontext + + import torch + from cuda.bindings import driver + + if not isinstance(plan, LaunchPlan): + raise GraphCaptureUnsupported( + "only frozen LaunchPlans are graph-capturable; per-call routes re-execute host logic" + ) + launches = tuple(support_launches) + tuple(plan._launches) + for launch in launches: + if launch._mode not in ("regular", "cluster"): + raise GraphCaptureUnsupported( + f"launch mode {launch._mode!r} has no validated graph-capture path" + ) + + # Captured launches build graph nodes and do not execute, so loom's CUDA + # watchdog (present only when the in-repo runtime shares this process) + # must not record completion events for them: the event record would be + # captured as a foreign node and the poller would query a captured event. + watchdog = sys.modules.get("loom.runtime.cuda_watchdog") + suspend = getattr(watchdog, "suspend_tracking", None) + suspension = suspend() if callable(suspend) else nullcontext() + + capture_stream = torch.cuda.Stream(device=plan.torch_stream.device) + cu_capture = driver.CUstream(capture_stream.cuda_stream) + nodes = [] + with suspension: + (err,) = driver.cuStreamBeginCapture( + cu_capture, driver.CUstreamCaptureMode.CU_STREAM_CAPTURE_MODE_THREAD_LOCAL + ) + _check_cu(err, "cuStreamBeginCapture failed") + try: + for launch in launches: + launch.launch(stream=capture_stream, timeout_ms=None) + info = driver.cuStreamGetCaptureInfo(cu_capture) + _check_cu(info[0], "cuStreamGetCaptureInfo failed") + status, leaves = info[1], info[4] + if ( + status != driver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_ACTIVE + or len(leaves) != 1 + ): + raise RuntimeError( + "graph capture did not add exactly one leaf node for a prepared launch" + ) + nodes.append(leaves[0]) + except BaseException: + driver.cuStreamEndCapture(cu_capture) # abandon the partial capture + raise + err, graph = driver.cuStreamEndCapture(cu_capture) + _check_cu(err, "cuStreamEndCapture failed") + + try: + err, _probe, total_nodes = driver.cuGraphGetNodes(graph, 0) + _check_cu(err, "cuGraphGetNodes failed") + if int(total_nodes) != len(launches): + raise RuntimeError( + f"captured graph has {int(total_nodes)} nodes, expected {len(launches)}; " + "foreign work was injected into the capture" + ) + for node in nodes: + err, node_type = driver.cuGraphNodeGetType(node) + _check_cu(err, "cuGraphNodeGetType failed") + if node_type != driver.CUgraphNodeType.CU_GRAPH_NODE_TYPE_KERNEL: + raise RuntimeError("captured graph node is not a kernel node") + err, graph_exec = driver.cuGraphInstantiate(graph, 0) + _check_cu(err, "cuGraphInstantiate failed") + except BaseException: + driver.cuGraphDestroy(graph) + raise + + node_params = [] + for launch, node in zip(launches, nodes): + params = driver.CUDA_KERNEL_NODE_PARAMS() + params.func = launch._kernel._func + params.gridDimX, params.gridDimY, params.gridDimZ = launch._grid + params.blockDimX, params.blockDimY, params.blockDimZ = launch._block + params.sharedMemBytes = launch._shared_mem + params.kernelParams = ctypes.addressof(launch._packed) + params.extra = 0 + node_params.append((node, params)) + return GraphExecPlan( + plan, + launches, + graph, + graph_exec, + node_params, + driver.CUstream(plan.stream_handle), + ) diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_row_norm.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/_row_norm.cu new file mode 100644 index 00000000..fc69f7b1 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_row_norm.cu @@ -0,0 +1,56 @@ +#include + +extern "C" __global__ void flashlib_cake_bf16_pair_row_norm( + const __nv_bfloat16* x, + const __nv_bfloat16* centroids, + float* x_sq, + float* c_sq, + long long x_rows, + long long c_rows, + int dim, + int compute_x, + int compute_c) { + const long long row = static_cast(blockIdx.x); + const bool is_x = row < x_rows; + if ((is_x && compute_x == 0) || (!is_x && compute_c == 0)) { + return; + } + + const long long local_row = is_x ? row : row - x_rows; + if (!is_x && local_row >= c_rows) { + return; + } + const __nv_bfloat16* input = is_x ? x : centroids; + float* output = is_x ? x_sq : c_sq; + const long long row_offset = local_row * static_cast(dim); + + float sum = 0.0f; + for (int column = static_cast(threadIdx.x); column < dim; + column += static_cast(blockDim.x)) { + const float value = __bfloat162float(input[row_offset + column]); + sum += value * value; + } + + constexpr unsigned int kFullMask = 0xffffffffu; + for (int offset = 16; offset > 0; offset >>= 1) { + sum += __shfl_down_sync(kFullMask, sum, offset); + } + + __shared__ float warp_sums[8]; + const int lane = static_cast(threadIdx.x) & 31; + const int warp = static_cast(threadIdx.x) >> 5; + if (lane == 0) { + warp_sums[warp] = sum; + } + __syncthreads(); + + if (warp == 0) { + sum = lane < (blockDim.x >> 5) ? warp_sums[lane] : 0.0f; + for (int offset = 16; offset > 0; offset >>= 1) { + sum += __shfl_down_sync(kFullMask, sum, offset); + } + if (lane == 0) { + output[local_row] = sum; + } + } +} diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_row_norm.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_row_norm.py new file mode 100644 index 00000000..eea6567f --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_row_norm.py @@ -0,0 +1,218 @@ +"""Standalone fused BF16 pair row norms; minimum architecture: sm_80.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from .kernels import ExportedKernel, KernelSpec + +_PAIR_ROW_NORM_KERNEL = ExportedKernel( + KernelSpec( + name="flashlib_cake_bf16_pair_row_norm", + symbol="flashlib_cake_bf16_pair_row_norm", + source="_row_norm.cu", + threads=256, + shared_mem_bytes=0, + cluster_dims=(1, 1, 1), + launch_mode="standard", + parameters=( + {"name": "x", "ctype": "const __nv_bfloat16 *"}, + {"name": "centroids", "ctype": "const __nv_bfloat16 *"}, + {"name": "x_sq", "ctype": "float *"}, + {"name": "c_sq", "ctype": "float *"}, + {"name": "x_rows", "ctype": "int64_t"}, + {"name": "c_rows", "ctype": "int64_t"}, + {"name": "dim", "ctype": "int32_t"}, + {"name": "compute_x", "ctype": "int32_t"}, + {"name": "compute_c", "ctype": "int32_t"}, + ), + specializations={}, + compile_options=("--std=c++17",), + ) +) + + +def _validate( + x: Any, + centroids: Any, + x_sq: Any, + c_sq: Any, +) -> tuple[int, int, int]: + import torch + + if not all(isinstance(item, torch.Tensor) and item.is_cuda for item in (x, centroids)): + raise TypeError("pair-row-norm inputs must be CUDA torch.Tensor objects") + if x.dtype is not torch.bfloat16 or centroids.dtype is not torch.bfloat16: + raise TypeError("pair-row-norm inputs must have bfloat16 dtype") + if x.ndim != 3 or centroids.ndim != 3 or not x.is_contiguous() or not centroids.is_contiguous(): + raise ValueError("pair-row-norm inputs must be contiguous [B, rows, D] tensors") + bsz, x_rows, dim = map(int, x.shape) + c_bsz, c_rows, c_dim = map(int, centroids.shape) + if (bsz, dim) != (c_bsz, c_dim) or x.device != centroids.device: + raise ValueError("pair-row-norm inputs must have matching batch, feature, and device") + for name, output, shape in ( + ("x_sq", x_sq, (bsz, x_rows)), + ("c_sq", c_sq, (bsz, c_rows)), + ): + if not isinstance(output, torch.Tensor) or not output.is_cuda: + raise TypeError(f"{name} must be a CUDA torch.Tensor") + if tuple(output.shape) != shape or output.dtype is not torch.float32: + raise ValueError(f"{name} must have dtype float32 and shape {shape}") + if output.device != x.device or not output.is_contiguous(): + raise ValueError(f"{name} must be contiguous and on {x.device}") + return bsz * x_rows, bsz * c_rows, dim + + +def _threads_for_dim(dim: int) -> int: + return min(256, max(32, 1 << (int(dim) - 1).bit_length())) + + +@dataclass +class PreparedBF16PairRowNorm: + """One pointer-rebindable launch for either or both KMeans row norms.""" + + launch_plan: Any + x_rows: int + c_rows: int + dim: int + compute_x: bool + compute_c: bool + _input_carriers: Any = field(default=None, repr=False) + + def rebind( + self, + x: Any, + centroids: Any, + x_sq: Any, + c_sq: Any, + *, + stream: Any = None, + ) -> None: + x_rows, c_rows, dim = _validate(x, centroids, x_sq, c_sq) + if (x_rows, c_rows, dim) != (self.x_rows, self.c_rows, self.dim): + raise RuntimeError("pair-row-norm prepared launch topology changed") + self.launch_plan.rebind_arguments( + {0: x, 1: centroids, 2: x_sq, 3: c_sq}, + stream=stream, + ) + + def launch(self, *, stream: Any = None, timeout_ms: float | None = None) -> None: + self.launch_plan.launch(stream=stream, timeout_ms=timeout_ms) + + def bind_hot(self, x: Any, centroids: Any) -> None: + """Overwrite both input pointer carriers in place without submitting. + + Graph-captured runtimes bind here and replay the captured kernel + chain themselves; the carriers target the same persistent packed + argument buffer the captured node's parameter update reads. The + prepared norm outputs stay bound (slot-owned scratch with stable + pointers), so the caller's signature cache must guarantee the tensor + topology matches the preparation. + """ + + carriers = self._input_carriers + if carriers is None: + prevent_gc = self.launch_plan._packed._prevent_gc + carriers = (prevent_gc[0], prevent_gc[1]) + self._input_carriers = carriers + carriers[0].value = x.data_ptr() + carriers[1].value = centroids.data_ptr() + + def launch_hot(self, x: Any, centroids: Any) -> None: + """Overwrite both input pointer carriers and submit the launch. + + The launch goes to its preparation-time stream — no re-marshal, no + per-launch stream query. + """ + + self.bind_hot(x, centroids) + self.launch_plan.launch(stream=None, timeout_ms=None) + + def record_stream(self, stream: Any) -> None: + """Tie every prepared launch argument to ``stream`` before release.""" + + for value in self.launch_plan._keepalive: + record_stream = getattr(value, "record_stream", None) + if callable(record_stream): + record_stream(stream) + + def release_bound_callers(self, keepalive: Any, *, stream: Any = None) -> None: + """Drop caller tensors while retaining one slot-owned scratch tensor.""" + + self.launch_plan.rebind_arguments( + {0: keepalive, 1: keepalive, 2: keepalive, 3: keepalive}, + stream=stream, + ) + + +def prepare_bf16_pair_row_norm( + x: Any, + centroids: Any, + x_sq: Any, + c_sq: Any, + *, + compute_x: bool, + compute_c: bool, + arch: str | None = None, + stream: Any = None, +) -> PreparedBF16PairRowNorm: + if not compute_x and not compute_c: + raise ValueError("pair-row-norm preparation requires at least one output") + x_rows, c_rows, dim = _validate(x, centroids, x_sq, c_sq) + launch_plan = _PAIR_ROW_NORM_KERNEL.prepare_launch( + x, + centroids, + x_sq, + c_sq, + x_rows, + c_rows, + dim, + int(compute_x), + int(compute_c), + grid=(x_rows + c_rows, 1, 1), + block=(_threads_for_dim(dim), 1, 1), + stream=stream, + arch=arch, + ) + return PreparedBF16PairRowNorm( + launch_plan=launch_plan, + x_rows=x_rows, + c_rows=c_rows, + dim=dim, + compute_x=bool(compute_x), + compute_c=bool(compute_c), + ) + + +def launch_bf16_pair_row_norm( + x: Any, + centroids: Any, + x_sq: Any, + c_sq: Any, + *, + compute_x: bool, + compute_c: bool, + stream: Any, + arch: str, + timeout_ms: float | None, +) -> None: + """Compute either or both BF16 row-squared norms in one CUDA activity.""" + + prepared = prepare_bf16_pair_row_norm( + x, + centroids, + x_sq, + c_sq, + compute_x=compute_x, + compute_c=compute_c, + arch=arch, + stream=stream, + ) + try: + prepared.launch(stream=stream, timeout_ms=timeout_ms) + finally: + prepared.release_bound_callers( + x_sq if compute_x else c_sq, + stream=stream, + ) diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/_runtime.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/_runtime.py new file mode 100644 index 00000000..64de997a --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/_runtime.py @@ -0,0 +1,912 @@ +from __future__ import annotations + +import contextlib +import contextvars +import ctypes +import os +import shutil +import threading +import time +from collections.abc import Iterator, Sequence +from typing import Any + +import torch +from cuda.bindings import driver, nvrtc + + +_LAUNCH_DEFAULTS: contextvars.ContextVar[tuple[str | None, Any, float | None] | None] = ( + contextvars.ContextVar("loom_export_launch_defaults", default=None) +) +_COMPILATION_CACHE_LOCK = threading.RLock() +_CUBIN_CACHE: dict[tuple[str, str, tuple[str, ...]], bytes] = {} +_MODULE_CACHE: dict[tuple[str, str, tuple[str, ...], int], "_LoadedModule"] = {} +_KERNEL_CACHE: dict[tuple[str, str, tuple[str, ...], int, str], "CUDAKernel"] = {} +_COMPILATION_CACHE_GENERATION = 0 +_RUNTIME_ACTIVITY_COUNTS = { + "source_reads": 0, + "nvrtc_compiles": 0, + "module_loads": 0, +} + + +def _record_runtime_activity(name: str) -> None: + with _COMPILATION_CACHE_LOCK: + _RUNTIME_ACTIVITY_COUNTS[name] += 1 + + +def record_source_read() -> None: + _record_runtime_activity("source_reads") + + +def runtime_activity_snapshot() -> dict[str, int]: + with _COMPILATION_CACHE_LOCK: + return dict(_RUNTIME_ACTIVITY_COUNTS) + + +class NVRTCError(RuntimeError): + def __init__(self, message: str, log: str): + self.log = log + super().__init__(f"{message}\n--- NVRTC log ---\n{log}") + + +def _check(err: int, msg: str) -> None: + if err != 0: + raise RuntimeError(f"{msg}: CUresult={err}") + + +def _nvrtc_check(err: int, msg: str) -> None: + if err != 0: + raise RuntimeError(f"{msg}: nvrtcResult={err}") + + +def _arch_flag_for_cc(major: int, minor: int) -> str: + sm = int(major) * 10 + int(minor) + return f"sm_{sm}a" if sm >= 90 else f"sm_{sm}" + + +def detect_gpu_arch() -> str: + forced = os.environ.get("LOOM_EXPORTED_FORCE_ARCH") or os.environ.get("LOOM_FORCE_ARCH") + if forced: + return forced + try: + (err,) = driver.cuInit(0) + if err != 0: + return "sm_100a" + err, dev = driver.cuDeviceGet(0) + if err != 0: + return "sm_100a" + err, major = driver.cuDeviceGetAttribute( + driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, dev + ) + err2, minor = driver.cuDeviceGetAttribute( + driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, dev + ) + if err != 0 or err2 != 0: + return "sm_100a" + return _arch_flag_for_cc(int(major), int(minor)) + except Exception: + return "sm_100a" + + +def resolve_gpu_arch(arch: str | None) -> str: + return detect_gpu_arch() if arch is None else str(arch) + + +def current_cuda_device_index() -> int: + return int(torch.cuda.current_device()) + + +def _is_torch_stream(stream: Any) -> bool: + return isinstance(stream, torch.cuda.Stream) + + +def launch_stream_context(stream: Any): + return torch.cuda.stream(stream) if _is_torch_stream(stream) else contextlib.nullcontext() + + +@contextlib.contextmanager +def launch_context( + *, + arch: str | None = None, + stream: Any = None, + timeout_ms: float | None = None, +) -> Iterator[None]: + '''Apply semantic-call launch defaults to every frozen kernel stage. + + The current PyTorch stream is captured lazily once before the first frozen + launch when no explicit stream is supplied. Passing a PyTorch stream also + makes dispatcher-side PyTorch tensor operations use that stream, not only + the generated driver launches. Context variables keep concurrent threads + and nested semantic calls isolated. + ''' + resolved_stream = stream + token = _LAUNCH_DEFAULTS.set((arch, resolved_stream, timeout_ms)) + stream_context = launch_stream_context(resolved_stream) + try: + with stream_context: + yield + finally: + _LAUNCH_DEFAULTS.reset(token) + + +def resolve_launch_defaults( + *, + arch: str | None, + stream: Any, + timeout_ms: float | None, +) -> tuple[str | None, Any, float | None]: + defaults = _LAUNCH_DEFAULTS.get() + if defaults is None: + return arch, stream, timeout_ms + default_arch, default_stream, default_timeout_ms = defaults + if stream is None and default_stream is None: + default_stream = torch.cuda.current_stream() + _LAUNCH_DEFAULTS.set((default_arch, default_stream, default_timeout_ms)) + return ( + default_arch if arch is None else arch, + default_stream if stream is None else stream, + default_timeout_ms if timeout_ms is None else timeout_ms, + ) + + +def compilation_cache_generation() -> int: + return _COMPILATION_CACHE_GENERATION + + +def _cuda_include_dirs() -> list[str]: + candidates = ["/usr/local/cuda/include"] + nvcc = shutil.which("nvcc") + if nvcc: + candidates.insert(0, os.path.join(os.path.dirname(os.path.dirname(nvcc)), "include")) + return [d for d in candidates if os.path.isdir(d)] + + +def _get_compile_log(prog: nvrtc.nvrtcProgram) -> str: + err, size = nvrtc.nvrtcGetProgramLogSize(prog) + if err != 0 or size <= 1: + return "" + log = b"\x00" * size + nvrtc.nvrtcGetProgramLog(prog, log) + return log.decode(errors="replace").rstrip("\x00") + + +def compile_cuda( + source: str, + *, + arch: str | None = None, + name: str = "kernel.cu", + options: list[str] | None = None, +) -> bytes: + _record_runtime_activity("nvrtc_compiles") + if arch is None: + arch = detect_gpu_arch() + opts = [ + f"--gpu-architecture={arch}", + "-std=c++17", + "-default-device", + ] + for include_dir in _cuda_include_dirs(): + opts.append(f"-I{include_dir}") + cccl = os.path.join(include_dir, "cccl") + if os.path.exists(os.path.join(cccl, "cuda", "std")): + opts.append(f"-I{cccl}") + if options: + opts.extend(options) + + err, prog = nvrtc.nvrtcCreateProgram(source.encode(), name.encode(), 0, [], []) + _nvrtc_check(err, "nvrtcCreateProgram failed") + try: + encoded_opts = [opt.encode() for opt in opts] + (err,) = nvrtc.nvrtcCompileProgram(prog, len(encoded_opts), encoded_opts) + if err != 0: + raise NVRTCError(f"Compilation failed for {name!r}", _get_compile_log(prog)) + err, size = nvrtc.nvrtcGetCUBINSize(prog) + _nvrtc_check(err, "nvrtcGetCUBINSize failed") + cubin = b"\x00" * size + (err,) = nvrtc.nvrtcGetCUBIN(prog, cubin) + _nvrtc_check(err, "nvrtcGetCUBIN failed") + return cubin + finally: + nvrtc.nvrtcDestroyProgram(prog) + + +def _is_tensor(arg: Any) -> bool: + return isinstance(arg, torch.Tensor) + + +def _marshal_arg(arg: Any, ctype: str) -> ctypes._SimpleCData: + normalized = " ".join(ctype.replace("__restrict__", "").split()) + if "*" in normalized: + if _is_tensor(arg): + return ctypes.c_void_p(arg.data_ptr()) + if isinstance(arg, ctypes.c_void_p): + return arg + if isinstance(arg, int) and not isinstance(arg, bool): + return ctypes.c_void_p(arg) + raise TypeError(f"Pointer argument {ctype!r} requires a CUDA tensor or integer device pointer, got {type(arg)}") + if normalized in {"bool", "_Bool"}: + if not isinstance(arg, bool): + raise TypeError(f"Scalar argument {ctype!r} requires bool, got {type(arg)}") + return ctypes.c_bool(arg) + integer_abis = { + "int8_t": (ctypes.c_int8, -(1 << 7), (1 << 7) - 1), + "signed char": (ctypes.c_int8, -(1 << 7), (1 << 7) - 1), + "uint8_t": (ctypes.c_uint8, 0, (1 << 8) - 1), + "unsigned char": (ctypes.c_uint8, 0, (1 << 8) - 1), + "int16_t": (ctypes.c_int16, -(1 << 15), (1 << 15) - 1), + "short": (ctypes.c_int16, -(1 << 15), (1 << 15) - 1), + "short int": (ctypes.c_int16, -(1 << 15), (1 << 15) - 1), + "uint16_t": (ctypes.c_uint16, 0, (1 << 16) - 1), + "unsigned short": (ctypes.c_uint16, 0, (1 << 16) - 1), + "unsigned short int": (ctypes.c_uint16, 0, (1 << 16) - 1), + "int32_t": (ctypes.c_int32, -(1 << 31), (1 << 31) - 1), + "int": (ctypes.c_int32, -(1 << 31), (1 << 31) - 1), + "signed int": (ctypes.c_int32, -(1 << 31), (1 << 31) - 1), + "uint32_t": (ctypes.c_uint32, 0, (1 << 32) - 1), + "unsigned int": (ctypes.c_uint32, 0, (1 << 32) - 1), + "int64_t": (ctypes.c_int64, -(1 << 63), (1 << 63) - 1), + "long long": (ctypes.c_int64, -(1 << 63), (1 << 63) - 1), + "uint64_t": (ctypes.c_uint64, 0, (1 << 64) - 1), + "unsigned long long": (ctypes.c_uint64, 0, (1 << 64) - 1), + "size_t": (ctypes.c_uint64, 0, (1 << 64) - 1), + } + integer_abi = integer_abis.get(normalized) + if integer_abi is not None: + if not isinstance(arg, int) or isinstance(arg, bool): + raise TypeError(f"Scalar argument {ctype!r} requires int, got {type(arg)}") + scalar_ctype, minimum, maximum = integer_abi + if arg < minimum or arg > maximum: + raise OverflowError(f"Scalar argument {ctype!r} is out of range: {arg}") + return scalar_ctype(arg) + if normalized in {"float", "double"}: + if not isinstance(arg, (int, float)) or isinstance(arg, bool): + raise TypeError(f"Scalar argument {ctype!r} requires int or float, got {type(arg)}") + return ctypes.c_float(arg) if normalized == "float" else ctypes.c_double(arg) + raise TypeError(f"Unsupported scalar ABI type {ctype!r}") + + +def _pack_args(args: Sequence[Any], arg_types: Sequence[str]) -> ctypes.Array: + if len(args) != len(arg_types): + raise ValueError(f"Argument count mismatch: got {len(args)}, expected {len(arg_types)}") + c_args = [_marshal_arg(arg, ctype) for arg, ctype in zip(args, arg_types, strict=True)] + ptrs = (ctypes.c_void_p * len(c_args))(*(ctypes.cast(ctypes.pointer(arg), ctypes.c_void_p) for arg in c_args)) + ptrs._prevent_gc = c_args # type: ignore[attr-defined] + return ptrs + + +class PreparedCUDAKernelLaunch: + # A fully marshalled generated-runtime launch reusable on the hot path. + + def __init__( + self, + kernel, + *, + mode, + grid, + block, + arg_types, + packed, + keepalive, + shared_mem, + cu_stream, + cluster_dims=None, + config=None, + ): + self._kernel = kernel + self._mode = mode + self._grid = tuple(grid) + self._block = tuple(block) + self._arg_types = tuple(arg_types) + self._packed = packed + self._keepalive = keepalive + self._shared_mem = shared_mem + self._cu_stream = cu_stream + self._cluster_dims = None if cluster_dims is None else tuple(cluster_dims) + self._config = config + + def rebind( + self, + kernel, + *, + mode, + grid, + block, + args, + arg_types, + shared_mem, + stream=None, + cluster_dims=None, + ): + # Reuse the existing void** and scalar carriers while replacing values. + + candidate_grid = tuple(grid) + candidate_block = tuple(block) + candidate_arg_types = tuple(arg_types) + candidate_cluster_dims = None if cluster_dims is None else tuple(cluster_dims) + mismatches = [] + if kernel is not self._kernel: + mismatches.append("kernel") + if mode != self._mode: + mismatches.append(f"mode ({self._mode!r} != {mode!r})") + if candidate_grid != self._grid: + mismatches.append(f"grid ({self._grid!r} != {candidate_grid!r})") + if candidate_block != self._block: + mismatches.append(f"block ({self._block!r} != {candidate_block!r})") + if int(shared_mem) != self._shared_mem: + mismatches.append(f"shared_mem ({self._shared_mem!r} != {int(shared_mem)!r})") + if candidate_cluster_dims != self._cluster_dims: + mismatches.append( + f"cluster_dims ({self._cluster_dims!r} != {candidate_cluster_dims!r})" + ) + if candidate_arg_types != self._arg_types: + mismatches.append("arg_types") + if mismatches: + raise RuntimeError( + "prepared CUDA launch topology mismatch: " + ", ".join(mismatches) + ) + + if len(args) != len(self._arg_types): + raise RuntimeError( + f"prepared CUDA launch argument count mismatch: " + f"expected {len(self._arg_types)}, got {len(args)}" + ) + return self.rebind_arguments(dict(enumerate(args)), stream=stream) + + def rebind_arguments(self, replacements, *, stream=None): + # Update selected ABI carriers without rebuilding or traversing the launch. + + old_c_args = self._packed._prevent_gc + if len(old_c_args) != len(self._arg_types): + raise RuntimeError("prepared CUDA launch has a corrupt packed argument array") + rebound = [] + for index, arg in replacements.items(): + if type(index) is not int or index < 0 or index >= len(self._arg_types): + raise IndexError(f"prepared CUDA launch argument index is out of range: {index!r}") + new_arg = _marshal_arg(arg, self._arg_types[index]) + old_arg = old_c_args[index] + if type(old_arg) is not type(new_arg): + raise RuntimeError( + f"prepared CUDA launch ABI mismatch at argument {index}: " + f"{type(old_arg).__name__} != {type(new_arg).__name__}" + ) + rebound.append((index, arg, old_arg, new_arg)) + + cu_stream = self._kernel._cu_stream(stream) + keepalive = list(self._keepalive) + for index, arg, old_arg, new_arg in rebound: + old_arg.value = new_arg.value + keepalive[index] = arg + self._keepalive = tuple(keepalive) + self._cu_stream = cu_stream + return self + + def rebind_tensor_arguments( + self, + bindings, + inputs, + *, + stream=None, + preserve_stream=False, + retain_inputs=True, + pointer_values=None, + inputs_already_scrubbed=False, + ): + # Captured public bindings are pointer-only. Update their existing + # c_void_p carriers directly instead of rematerializing ctypes objects. + + if not isinstance(preserve_stream, bool): + raise TypeError("preserve_stream must be a bool") + if not isinstance(retain_inputs, bool): + raise TypeError("retain_inputs must be a bool") + if not isinstance(inputs_already_scrubbed, bool): + raise TypeError("inputs_already_scrubbed must be a bool") + if inputs_already_scrubbed and retain_inputs: + raise ValueError("inputs_already_scrubbed requires retain_inputs=False") + if pointer_values is not None and retain_inputs: + raise ValueError("pointer_values requires retain_inputs=False") + old_c_args = self._packed._prevent_gc + if len(old_c_args) != len(self._arg_types): + raise RuntimeError("prepared CUDA launch has a corrupt packed argument array") + replacements = [] + for index, key in bindings: + if type(index) is not int or index < 0 or index >= len(self._arg_types): + raise IndexError(f"prepared CUDA launch argument index is out of range: {index!r}") + arg = None + if pointer_values is None: + try: + arg = inputs[key] + except KeyError: + raise KeyError(f"missing prepared CUDA tensor binding: {key!r}") from None + data_ptr = getattr(arg, "data_ptr", None) + if not callable(data_ptr): + raise TypeError(f"prepared CUDA tensor binding {key!r} is not tensor-like") + pointer = int(data_ptr()) + else: + try: + pointer = pointer_values[key] + except KeyError: + raise KeyError(f"missing prepared CUDA pointer binding: {key!r}") from None + if type(pointer) is not int: + raise TypeError(f"prepared CUDA pointer binding {key!r} is not an int") + old_arg = old_c_args[index] + if type(old_arg) is not ctypes.c_void_p: + raise RuntimeError( + f"prepared CUDA tensor binding {key!r} does not target a pointer carrier" + ) + if inputs_already_scrubbed and type(self._keepalive[index]) is not int: + raise RuntimeError( + f"prepared CUDA tensor binding {key!r} retained an unexpected caller reference" + ) + replacements.append((index, key, arg, pointer, old_arg)) + + keepalive = None if inputs_already_scrubbed else list(self._keepalive) + for index, key, arg, pointer, old_arg in replacements: + old_arg.value = pointer + if keepalive is not None: + if retain_inputs and arg is None: + arg = inputs[key] + keepalive[index] = arg if retain_inputs else pointer + if keepalive is not None: + self._keepalive = tuple(keepalive) + if not preserve_stream: + self._cu_stream = self._kernel._cu_stream(stream) + return self + + def _scrub_stream_bound_pointer_keepalives(self, bindings): + '''Drop selected tensor references without changing carriers or stream.''' + + old_c_args = self._packed._prevent_gc + if len(old_c_args) != len(self._arg_types): + raise RuntimeError("prepared CUDA launch has a corrupt packed argument array") + scrubbed = list(self._keepalive) + program = [] + for index, key in bindings: + if type(index) is not int or index < 0 or index >= len(self._arg_types): + raise IndexError(f"prepared CUDA launch argument index is out of range: {index!r}") + carrier = old_c_args[index] + if type(carrier) is not ctypes.c_void_p: + raise RuntimeError( + f"prepared CUDA tensor binding {key!r} does not target a pointer carrier" + ) + value = scrubbed[index] + data_ptr = getattr(value, "data_ptr", None) + if not callable(data_ptr): + raise RuntimeError( + f"prepared CUDA tensor binding {key!r} retained an unexpected value" + ) + pointer = int(data_ptr()) + if carrier.value != pointer: + raise RuntimeError( + f"prepared CUDA tensor binding {key!r} carrier disagrees with its keepalive" + ) + scrubbed[index] = pointer + program.append((key, carrier)) + self._keepalive = tuple(scrubbed) + return tuple(program) + + def launch(self, *, stream=None, timeout_ms=None): + kernel = self._kernel + if kernel._closed: + raise RuntimeError("Kernel has been unloaded") + cu_stream = self._cu_stream if stream is None else kernel._cu_stream(stream) + if self._mode == "cluster": + self._config.hStream = cu_stream + (err,) = driver.cuLaunchKernelEx(self._config, kernel._func, self._packed, 0) + api = "cuLaunchKernelEx" + elif self._mode == "cooperative": + (err,) = driver.cuLaunchCooperativeKernel( + kernel._func, + *self._grid, + *self._block, + self._shared_mem, + cu_stream, + self._packed, + ) + api = "cuLaunchCooperativeKernel" + else: + (err,) = driver.cuLaunchKernel( + kernel._func, + *self._grid, + *self._block, + self._shared_mem, + cu_stream, + self._packed, + 0, + ) + api = "cuLaunchKernel" + _check(err, f"{api} failed for {kernel._func_name!r}") + if timeout_ms is not None: + kernel._wait_with_timeout(cu_stream, timeout_ms) + + +class _LoadedModule: + def __init__(self, cubin: bytes): + _record_runtime_activity("module_loads") + torch.empty(0, device="cuda") + err, self._module = driver.cuModuleLoadData(cubin) + _check(err, "cuModuleLoadData failed") + self._functions: dict[str, Any] = {} + self._closed = False + + @property + def handle(self): + return self._module + + @property + def closed(self) -> bool: + return self._closed + + def function(self, func_name: str): + if self._closed: + raise RuntimeError("CUDA module has been unloaded") + function = self._functions.get(func_name) + if function is None: + err, function = driver.cuModuleGetFunction(self._module, func_name.encode()) + _check(err, f"cuModuleGetFunction failed for {func_name!r}") + self._functions[func_name] = function + return function + + def close(self) -> None: + if not self._closed: + driver.cuModuleUnload(self._module) + self._closed = True + + +class CUDAKernel: + def __init__(self, cubin: bytes, func_name: str): + self._module_owner = _LoadedModule(cubin) + self._owns_module = True + self._initialize_from_module(func_name) + + @classmethod + def from_loaded_module(cls, module: _LoadedModule, func_name: str) -> "CUDAKernel": + kernel = cls.__new__(cls) + kernel._module_owner = module + kernel._owns_module = False + kernel._initialize_from_module(func_name) + return kernel + + def _initialize_from_module(self, func_name: str) -> None: + self._closed = True + self._func_name = func_name + self._module = self._module_owner.handle + self._func = self._module_owner.function(func_name) + self._dynamic_smem_opt_in_bytes = 0 + self._closed = False + + @property + def closed(self) -> bool: + return self._closed + + def _ensure_dynamic_smem_opt_in(self, shared_mem: int) -> None: + if shared_mem <= 48 * 1024 or shared_mem <= self._dynamic_smem_opt_in_bytes: + return + (err,) = driver.cuFuncSetAttribute( + self._func, + driver.CUfunction_attribute.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + shared_mem, + ) + _check(err, "cuFuncSetAttribute MAX_DYNAMIC_SHARED_SIZE_BYTES failed") + self._dynamic_smem_opt_in_bytes = max(self._dynamic_smem_opt_in_bytes, shared_mem) + + def _cu_stream(self, stream=None): + if stream is None: + stream = torch.cuda.current_stream() + handle = getattr(stream, "cuda_stream", stream) + return driver.CUstream(int(handle)) + + def _wait_with_timeout(self, cu_stream: driver.CUstream, timeout_ms: float) -> None: + deadline = time.monotonic() + timeout_ms / 1000.0 + while True: + (err,) = driver.cuStreamQuery(cu_stream) + if err == 0: + return + if err != 600: + _check(err, f"cuStreamQuery failed for {self._func_name!r}") + if time.monotonic() >= deadline: + raise TimeoutError(f"Kernel {self._func_name!r} did not complete within {timeout_ms:.0f} ms") + time.sleep(0.001) + + def launch( + self, + *, + grid: tuple[int, int, int], + block: tuple[int, int, int], + args: Sequence[Any], + arg_types: Sequence[str], + shared_mem: int = 0, + stream=None, + timeout_ms: float | None = None, + ) -> None: + self.prepare_launch( + grid=grid, + block=block, + args=args, + arg_types=arg_types, + shared_mem=shared_mem, + stream=stream, + ).launch(timeout_ms=timeout_ms) + + def prepare_launch( + self, + *, + grid, + block, + args, + arg_types, + shared_mem=0, + stream=None, + ): + if self._closed: + raise RuntimeError("Kernel has been unloaded") + self._ensure_dynamic_smem_opt_in(shared_mem) + return PreparedCUDAKernelLaunch( + self, + mode="regular", + grid=grid, + block=block, + arg_types=arg_types, + packed=_pack_args(args, arg_types), + keepalive=tuple(args), + shared_mem=shared_mem, + cu_stream=self._cu_stream(stream), + ) + + def rebind_launch( + self, + prepared, + *, + grid, + block, + args, + arg_types, + shared_mem=0, + stream=None, + ): + if not isinstance(prepared, PreparedCUDAKernelLaunch): + raise TypeError("prepared must be a PreparedCUDAKernelLaunch") + return prepared.rebind( + self, + mode="regular", + grid=grid, + block=block, + args=args, + arg_types=arg_types, + shared_mem=shared_mem, + stream=stream, + ) + + def launch_cluster( + self, + *, + grid: tuple[int, int, int], + block: tuple[int, int, int], + args: Sequence[Any], + arg_types: Sequence[str], + cluster_dims: tuple[int, int, int], + shared_mem: int = 0, + stream=None, + timeout_ms: float | None = None, + ) -> None: + self.prepare_launch_cluster( + grid=grid, + block=block, + args=args, + arg_types=arg_types, + cluster_dims=cluster_dims, + shared_mem=shared_mem, + stream=stream, + ).launch(timeout_ms=timeout_ms) + + def prepare_launch_cluster( + self, + *, + grid, + block, + args, + arg_types, + cluster_dims, + shared_mem=0, + stream=None, + ): + if self._closed: + raise RuntimeError("Kernel has been unloaded") + self._ensure_dynamic_smem_opt_in(shared_mem) + packed = _pack_args(args, arg_types) + cu_stream = self._cu_stream(stream) + + attr_cluster = driver.CUlaunchAttribute() + attr_cluster.id = driver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION + attr_cluster.value.clusterDim.x = cluster_dims[0] + attr_cluster.value.clusterDim.y = cluster_dims[1] + attr_cluster.value.clusterDim.z = cluster_dims[2] + + attr_sched = driver.CUlaunchAttribute() + attr_sched.id = driver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE + attr_sched.value.clusterSchedulingPolicyPreference = ( + driver.CUclusterSchedulingPolicy.CU_CLUSTER_SCHEDULING_POLICY_SPREAD + ) + + config = driver.CUlaunchConfig() + config.gridDimX = grid[0] + config.gridDimY = grid[1] + config.gridDimZ = grid[2] + config.blockDimX = block[0] + config.blockDimY = block[1] + config.blockDimZ = block[2] + config.sharedMemBytes = shared_mem + config.hStream = cu_stream + config.attrs = [attr_cluster, attr_sched] + config.numAttrs = 2 + return PreparedCUDAKernelLaunch( + self, + mode="cluster", + grid=grid, + block=block, + arg_types=arg_types, + packed=packed, + keepalive=tuple(args), + shared_mem=shared_mem, + cu_stream=cu_stream, + cluster_dims=cluster_dims, + config=config, + ) + + def rebind_launch_cluster( + self, + prepared, + *, + grid, + block, + args, + arg_types, + cluster_dims, + shared_mem=0, + stream=None, + ): + if not isinstance(prepared, PreparedCUDAKernelLaunch): + raise TypeError("prepared must be a PreparedCUDAKernelLaunch") + return prepared.rebind( + self, + mode="cluster", + grid=grid, + block=block, + args=args, + arg_types=arg_types, + cluster_dims=cluster_dims, + shared_mem=shared_mem, + stream=stream, + ) + + def launch_cooperative( + self, + *, + grid: tuple[int, int, int], + block: tuple[int, int, int], + args: Sequence[Any], + arg_types: Sequence[str], + shared_mem: int = 0, + stream=None, + timeout_ms: float | None = None, + ) -> None: + self.prepare_launch_cooperative( + grid=grid, + block=block, + args=args, + arg_types=arg_types, + shared_mem=shared_mem, + stream=stream, + ).launch(timeout_ms=timeout_ms) + + def prepare_launch_cooperative( + self, + *, + grid, + block, + args, + arg_types, + shared_mem=0, + stream=None, + ): + if self._closed: + raise RuntimeError("Kernel has been unloaded") + self._ensure_dynamic_smem_opt_in(shared_mem) + return PreparedCUDAKernelLaunch( + self, + mode="cooperative", + grid=grid, + block=block, + arg_types=arg_types, + packed=_pack_args(args, arg_types), + keepalive=tuple(args), + shared_mem=shared_mem, + cu_stream=self._cu_stream(stream), + ) + + def rebind_launch_cooperative( + self, + prepared, + *, + grid, + block, + args, + arg_types, + shared_mem=0, + stream=None, + ): + if not isinstance(prepared, PreparedCUDAKernelLaunch): + raise TypeError("prepared must be a PreparedCUDAKernelLaunch") + return prepared.rebind( + self, + mode="cooperative", + grid=grid, + block=block, + args=args, + arg_types=arg_types, + shared_mem=shared_mem, + stream=stream, + ) + + def close(self) -> None: + if not self._closed: + if self._owns_module: + self._module_owner.close() + self._closed = True + + def __del__(self): + try: + self.close() + except Exception: + pass + + +def load_cached_kernel( + source: str, + *, + source_digest: str, + func_name: str, + arch: str, + device_index: int, + name: str, + options: Sequence[str] = (), +) -> CUDAKernel: + '''Compile/load one content-addressed kernel shared by equivalent aliases. + + Identical source, architecture, and options share NVRTC output. The active + device is added for module loading, and the function symbol only for the + final wrapper, so a multi-entrypoint translation unit loads one module per + device without conflating its functions. + ''' + option_tuple = tuple(options) + cubin_key = (source_digest, arch, option_tuple) + module_key = (*cubin_key, int(device_index)) + kernel_key = (*module_key, func_name) + with _COMPILATION_CACHE_LOCK: + kernel = _KERNEL_CACHE.get(kernel_key) + if kernel is not None and not kernel.closed: + return kernel + cubin = _CUBIN_CACHE.get(cubin_key) + if cubin is None: + cubin = compile_cuda(source, arch=arch, name=name, options=list(option_tuple)) + _CUBIN_CACHE[cubin_key] = cubin + module = _MODULE_CACHE.get(module_key) + if module is None or module.closed: + module = _LoadedModule(cubin) + _MODULE_CACHE[module_key] = module + kernel = CUDAKernel.from_loaded_module(module, func_name) + _KERNEL_CACHE[kernel_key] = kernel + return kernel + + +def clear_compilation_cache() -> None: + '''Clear process-local cubin/module caches (primarily for diagnostics).''' + global _COMPILATION_CACHE_GENERATION + with _COMPILATION_CACHE_LOCK: + for kernel in _KERNEL_CACHE.values(): + kernel._closed = True + for module in _MODULE_CACHE.values(): + module.close() + _KERNEL_CACHE.clear() + _MODULE_CACHE.clear() + _CUBIN_CACHE.clear() + _COMPILATION_CACHE_GENERATION += 1 diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0000.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0000.cu new file mode 100644 index 00000000..df9612f5 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0000.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_lowdim_pack_e50c_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 16 * 8; + int row = vec_idx / 16; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 16 * 8; + int row_1 = vec_idx_1 / 16; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0001.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0001.cu new file mode 100644 index 00000000..bd103c07 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0001.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_lowdim_e50c_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0002.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0002.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0002.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0003.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0003.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0003.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0004.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0004.cu new file mode 100644 index 00000000..482ca687 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0004.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 24 * 8; + int row = vec_idx / 24; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 24 * 8; + int row_1 = vec_idx_1 / 24; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0005.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0005.cu new file mode 100644 index 00000000..7070fa90 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0005.cu @@ -0,0 +1,600 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 49152 +#define SMEM_SMEM_X_STRIDE 49152 +#define SMEM_SMEM_C_OFF 50176 +#define SMEM_SMEM_C_STAGE_BYTES 98304 +#define SMEM_SMEM_C_STRIDE 98304 +#define SMEM_SMEM_CSQ_OFF 148480 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 149504 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int smem_c_addr = smem + 50176; + float* smem_csq = reinterpret_cast(smem_raw + 148480); + const int smem_csq_addr = smem + 148480; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 49152); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 98304); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0006.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0006.cu new file mode 100644 index 00000000..482ca687 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0006.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 24 * 8; + int row = vec_idx / 24; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 24 * 8; + int row_1 = vec_idx_1 / 24; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0007.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0007.cu new file mode 100644 index 00000000..7070fa90 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0007.cu @@ -0,0 +1,600 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 49152 +#define SMEM_SMEM_X_STRIDE 49152 +#define SMEM_SMEM_C_OFF 50176 +#define SMEM_SMEM_C_STAGE_BYTES 98304 +#define SMEM_SMEM_C_STRIDE 98304 +#define SMEM_SMEM_CSQ_OFF 148480 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 149504 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int smem_c_addr = smem + 50176; + float* smem_csq = reinterpret_cast(smem_raw + 148480); + const int smem_csq_addr = smem + 148480; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 49152); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 98304); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0008.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0008.cu new file mode 100644 index 00000000..17caae28 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0008.cu @@ -0,0 +1,1211 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 49152 +#define SMEM_SMEM_X0_STRIDE 49152 +#define SMEM_SMEM_X1_OFF 50176 +#define SMEM_SMEM_X1_STAGE_BYTES 49152 +#define SMEM_SMEM_X1_STRIDE 49152 +#define SMEM_SMEM_C_OFF 99328 +#define SMEM_SMEM_C_STAGE_BYTES 98304 +#define SMEM_SMEM_C_STRIDE 98304 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int smem_x1_addr = smem + 50176; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); + const int smem_c_addr = smem + 99328; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 49152); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 49152); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 98304); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x0_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_4 = smem_x1_addr + 16384; + int _mma_a_lo_4 = make_warp_uniform((_mma_a_addr_4 >> 4) & 0x3FFF); + int _mma_b_addr_4 = smem_c_addr + 32768; + int _mma_b_lo_4 = make_warp_uniform((_mma_b_addr_4 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_5 = smem_x1_addr + 32768; + int _mma_a_lo_5 = make_warp_uniform((_mma_a_addr_5 >> 4) & 0x3FFF); + int _mma_b_addr_5 = smem_c_addr + 65536; + int _mma_b_lo_5 = make_warp_uniform((_mma_b_addr_5 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0009.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0009.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0009.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0010.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0010.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0010.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0011.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0011.cu new file mode 100644 index 00000000..f6311654 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0011.cu @@ -0,0 +1,535 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 41984 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 43008 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitk_partial_blockn64_g2r4_b5a6_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, float* __restrict__ partial_scores, int* __restrict__ partial_indices, int B, int N, int D, int K, int num_n_tiles, int K_tiles, int K_slices) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 41984); + const int smem_csq_addr = smem + 41984; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: compute ---- + if (warp <= 3) { + { // compute_main + int total_work = B * num_n_tiles * K_slices; + int compute_warp = warp; + int lane_pair = lane % 4; + int row_origin = compute_warp * 16; + int row_lane_base = row_origin + lane / 4; + int row0 = row_lane_base; + int row1 = row_lane_base + 8; + int compute_tid = compute_warp * 32 + lane; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_slice = work_idx % (unsigned int)K_slices; + int point_tile_idx = work_idx / (unsigned int)K_slices; + int batch = point_tile_idx / num_n_tiles; + int slice_k_start = iter_slice * 2; + int csq_smem_addr = smem_csq_addr; + float best0 = -3.4e+38f; + float best1 = -3.4e+38f; + int idx0 = slice_k_start * 256; + int idx1 = idx0; + #pragma unroll 1 + for (int local_k = 0; local_k < 2; local_k++) { + int iter_k = slice_k_start + local_k; + int off_k = iter_k * 256; + if ((float)compute_tid < 64.0f) { + int csq_base = compute_tid * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch * K + off_k + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + #pragma unroll + for (int score_base = 0; score_base < 256; score_base += 128) { + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, [%64];" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])), "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) + : "r"(taddr + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int rep = 0; rep < 16.0f; rep++) { + int local_reg = rep * 4; + int col_base = score_base + rep * 8 + lane_pair * 2; + float csq0 = smem_csq[col_base]; + float csq1 = smem_csq[col_base + 1]; + float d0 = _tmem_load_0[local_reg] - csq0; + if (d0 > best0) { + best0 = d0; + idx0 = off_k + col_base; + } + float d1 = _tmem_load_0[local_reg + 1] - csq1; + if (d1 > best0) { + best0 = d1; + idx0 = off_k + col_base + 1; + } + float d2 = _tmem_load_0[local_reg + 2] - csq0; + if (d2 > best1) { + best1 = d2; + idx1 = off_k + col_base; + } + float d3 = _tmem_load_0[local_reg + 3] - csq1; + if (d3 > best1) { + best1 = d3; + idx1 = off_k + col_base + 1; + } + } + } + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best0, 1); + float peer0 = _shfl_xor_0; + int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, idx0, 1); + int peer0_idx = _shfl_xor_1; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer0_idx; + } + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, best0, 2); + peer0 = _shfl_xor_2; + int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, idx0, 2); + peer0_idx = _shfl_xor_3; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer0_idx; + } + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, best1, 1); + float peer1 = _shfl_xor_4; + int _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, idx1, 1); + int peer1_idx = _shfl_xor_5; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer1_idx; + } + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, best1, 2); + peer1 = _shfl_xor_6; + int _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, idx1, 2); + peer1_idx = _shfl_xor_7; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer1_idx; + } + if (lane_pair == 0) { + int partial_offset0 = work_idx * 64 + (unsigned int)row0; + *((float*)(partial_scores + partial_offset0)) = best0; + *((int*)(partial_indices + partial_offset0)) = idx0; + int partial_offset1 = work_idx * 64 + (unsigned int)row1; + *((float*)(partial_scores + partial_offset1)) = best1; + *((int*)(partial_indices + partial_offset1)) = idx1; + } + } + } + // ---- Role: load ---- + } else if (warp == 4) { + { // load_main + int total_work_1 = B * num_n_tiles * K_slices; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_1; work_idx_1 += num_bids) { + int iter_slice_1 = work_idx_1 % (unsigned int)K_slices; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_slices; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int n_tile = point_tile_idx_1 % num_n_tiles; + int off_n = n_tile * 64; + int x_row = batch_1 * N + off_n; + int slice_k_start_1 = iter_slice_1 * 2; + #pragma unroll 1 + for (int local_k_1 = 0; local_k_1 < 2; local_k_1++) { + int iter_k_1 = slice_k_start_1 + local_k_1; + int off_k_1 = iter_k_1 * 256; + int c_row = batch_1 * K + off_k_1; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 5) { + { // mma_main + int total_work_2 = B * num_n_tiles * K_slices; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = bid; work_idx_2 < total_work_2; work_idx_2 += num_bids) { + int iter_slice_2 = work_idx_2 % (unsigned int)K_slices; + int slice_k_start_2 = iter_slice_2 * 2; + #pragma unroll 1 + for (int local_k_2 = 0; local_k_2 < 2; local_k_2++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 71304336;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0012.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0012.cu new file mode 100644 index 00000000..461424c3 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0012.cu @@ -0,0 +1,81 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_highd_splitk_reduce_blockn64_g2r4_b5a6_v1(float* __restrict__ partial_scores, int* __restrict__ partial_indices, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_slices) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid / 4; + int row_lane = tid % 4; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 64 + row; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_slice = row_lane; iter_slice < K_slices; iter_slice += 4) { + int partial_offset = (point_tile_idx * (unsigned int)K_slices + (unsigned int)iter_slice) * 64 + (unsigned int)row; + float score = partial_scores[partial_offset]; + int idx = partial_indices[partial_offset]; + if (score > best_score) { + best_score = score; + best_idx = idx; + } + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best_score, 1); + float peer_score = _shfl_xor_0; + int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, best_idx, 1); + int peer_idx = _shfl_xor_1; + if (peer_score > best_score) { + best_score = peer_score; + best_idx = peer_idx; + } + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, best_score, 2); + peer_score = _shfl_xor_2; + int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, best_idx, 2); + peer_idx = _shfl_xor_3; + if (peer_score > best_score) { + best_score = peer_score; + best_idx = peer_idx; + } + if (row_lane == 0) { + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = best_idx; + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0013.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0013.cu new file mode 100644 index 00000000..5b836885 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0013.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_microdim_pack_6cd2_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 8 * 8; + int row = vec_idx / 8; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 8 * 8; + int row_1 = vec_idx_1 / 8; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0014.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0014.cu new file mode 100644 index 00000000..1c18f40d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0014.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_microdim_6cd2_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0015.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0015.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0015.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0016.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0016.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0016.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0017.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0017.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0017.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0018.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0018.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0018.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0019.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0019.cu new file mode 100644 index 00000000..5b836885 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0019.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_microdim_pack_6cd2_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 8 * 8; + int row = vec_idx / 8; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 8 * 8; + int row_1 = vec_idx_1 / 8; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0020.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0020.cu new file mode 100644 index 00000000..1c18f40d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0020.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_microdim_6cd2_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0021.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0021.cu new file mode 100644 index 00000000..5b836885 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0021.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_microdim_pack_6cd2_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 8 * 8; + int row = vec_idx / 8; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 8 * 8; + int row_1 = vec_idx_1 / 8; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0022.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0022.cu new file mode 100644 index 00000000..1c18f40d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0022.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_microdim_6cd2_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0023.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0023.cu new file mode 100644 index 00000000..5b836885 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0023.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_microdim_pack_6cd2_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 8 * 8; + int row = vec_idx / 8; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 8 * 8; + int row_1 = vec_idx_1 / 8; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0024.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0024.cu new file mode 100644 index 00000000..1c18f40d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0024.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_microdim_6cd2_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0025.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0025.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0025.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0026.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0026.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0026.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0027.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0027.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0027.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0028.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0028.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0028.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0029.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0029.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0029.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0030.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0030.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0030.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0031.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0031.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0031.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0032.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0032.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0032.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0033.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0033.cu new file mode 100644 index 00000000..df9612f5 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0033.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_lowdim_pack_e50c_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 16 * 8; + int row = vec_idx / 16; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 16 * 8; + int row_1 = vec_idx_1 / 16; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0034.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0034.cu new file mode 100644 index 00000000..bd103c07 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0034.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_lowdim_e50c_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0035.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0035.cu new file mode 100644 index 00000000..df9612f5 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0035.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_lowdim_pack_e50c_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 16 * 8; + int row = vec_idx / 16; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 16 * 8; + int row_1 = vec_idx_1 / 16; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0036.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0036.cu new file mode 100644 index 00000000..bd103c07 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0036.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_lowdim_e50c_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0037.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0037.cu new file mode 100644 index 00000000..df9612f5 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0037.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_lowdim_pack_e50c_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 16 * 8; + int row = vec_idx / 16; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 16 * 8; + int row_1 = vec_idx_1 / 16; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0038.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0038.cu new file mode 100644 index 00000000..bd103c07 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0038.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_lowdim_e50c_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0039.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0039.cu new file mode 100644 index 00000000..df9612f5 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0039.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_lowdim_pack_e50c_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 16 * 8; + int row = vec_idx / 16; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 16 * 8; + int row_1 = vec_idx_1 / 16; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0040.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0040.cu new file mode 100644 index 00000000..bd103c07 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0040.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_lowdim_e50c_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0041.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0041.cu new file mode 100644 index 00000000..df9612f5 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0041.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_lowdim_pack_e50c_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 16 * 8; + int row = vec_idx / 16; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 16 * 8; + int row_1 = vec_idx_1 / 16; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0042.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0042.cu new file mode 100644 index 00000000..bd103c07 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0042.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_lowdim_e50c_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0043.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0043.cu new file mode 100644 index 00000000..df9612f5 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0043.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_lowdim_pack_e50c_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 16 * 8; + int row = vec_idx / 16; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 16 * 8; + int row_1 = vec_idx_1 / 16; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0044.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0044.cu new file mode 100644 index 00000000..bd103c07 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0044.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_lowdim_e50c_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0045.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0045.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0045.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0046.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0046.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0046.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0047.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0047.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0047.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0048.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0048.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0048.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0049.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0049.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0049.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0050.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0050.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0050.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0051.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0051.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0051.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0052.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0052.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0052.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0053.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0053.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0053.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0054.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0054.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0054.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0055.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0055.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0055.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0056.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0056.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0056.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0057.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0057.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0057.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0058.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0058.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0058.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0059.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0059.cu new file mode 100644 index 00000000..fcaac275 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0059.cu @@ -0,0 +1,535 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 41984 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 43008 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitk_partial_blockn64_g1r4_streamdep_r63_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, float* __restrict__ partial_scores, int* __restrict__ partial_indices, int B, int N, int D, int K, int num_n_tiles, int K_tiles, int K_slices) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 41984); + const int smem_csq_addr = smem + 41984; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: compute ---- + if (warp <= 3) { + { // compute_main + int total_work = B * num_n_tiles * K_slices; + int compute_warp = warp; + int lane_pair = lane % 4; + int row_origin = compute_warp * 16; + int row_lane_base = row_origin + lane / 4; + int row0 = row_lane_base; + int row1 = row_lane_base + 8; + int compute_tid = compute_warp * 32 + lane; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_slice = work_idx % (unsigned int)K_slices; + int point_tile_idx = work_idx / (unsigned int)K_slices; + int batch = point_tile_idx / num_n_tiles; + int slice_k_start = iter_slice; + int csq_smem_addr = smem_csq_addr; + float best0 = -3.4e+38f; + float best1 = -3.4e+38f; + int idx0 = slice_k_start * 256; + int idx1 = idx0; + #pragma unroll 1 + for (int local_k = 0; local_k < 1; local_k++) { + int iter_k = slice_k_start + local_k; + int off_k = iter_k * 256; + if ((float)compute_tid < 64.0f) { + int csq_base = compute_tid * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch * K + off_k + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + #pragma unroll + for (int score_base = 0; score_base < 256; score_base += 128) { + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, [%64];" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])), "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) + : "r"(taddr + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int rep = 0; rep < 16.0f; rep++) { + int local_reg = rep * 4; + int col_base = score_base + rep * 8 + lane_pair * 2; + float csq0 = smem_csq[col_base]; + float csq1 = smem_csq[col_base + 1]; + float d0 = _tmem_load_0[local_reg] - csq0; + if (d0 > best0) { + best0 = d0; + idx0 = off_k + col_base; + } + float d1 = _tmem_load_0[local_reg + 1] - csq1; + if (d1 > best0) { + best0 = d1; + idx0 = off_k + col_base + 1; + } + float d2 = _tmem_load_0[local_reg + 2] - csq0; + if (d2 > best1) { + best1 = d2; + idx1 = off_k + col_base; + } + float d3 = _tmem_load_0[local_reg + 3] - csq1; + if (d3 > best1) { + best1 = d3; + idx1 = off_k + col_base + 1; + } + } + } + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best0, 1); + float peer0 = _shfl_xor_0; + int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, idx0, 1); + int peer0_idx = _shfl_xor_1; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer0_idx; + } + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, best0, 2); + peer0 = _shfl_xor_2; + int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, idx0, 2); + peer0_idx = _shfl_xor_3; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer0_idx; + } + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, best1, 1); + float peer1 = _shfl_xor_4; + int _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, idx1, 1); + int peer1_idx = _shfl_xor_5; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer1_idx; + } + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, best1, 2); + peer1 = _shfl_xor_6; + int _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, idx1, 2); + peer1_idx = _shfl_xor_7; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer1_idx; + } + if (lane_pair == 0) { + int partial_offset0 = work_idx * 64 + (unsigned int)row0; + *((float*)(partial_scores + partial_offset0)) = best0; + *((int*)(partial_indices + partial_offset0)) = idx0; + int partial_offset1 = work_idx * 64 + (unsigned int)row1; + *((float*)(partial_scores + partial_offset1)) = best1; + *((int*)(partial_indices + partial_offset1)) = idx1; + } + } + } + // ---- Role: load ---- + } else if (warp == 4) { + { // load_main + int total_work_1 = B * num_n_tiles * K_slices; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_1; work_idx_1 += num_bids) { + int iter_slice_1 = work_idx_1 % (unsigned int)K_slices; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_slices; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int n_tile = point_tile_idx_1 % num_n_tiles; + int off_n = n_tile * 64; + int x_row = batch_1 * N + off_n; + int slice_k_start_1 = iter_slice_1; + #pragma unroll 1 + for (int local_k_1 = 0; local_k_1 < 1; local_k_1++) { + int iter_k_1 = slice_k_start_1 + local_k_1; + int off_k_1 = iter_k_1 * 256; + int c_row = batch_1 * K + off_k_1; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 5) { + { // mma_main + int total_work_2 = B * num_n_tiles * K_slices; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = bid; work_idx_2 < total_work_2; work_idx_2 += num_bids) { + int iter_slice_2 = work_idx_2 % (unsigned int)K_slices; + int slice_k_start_2 = iter_slice_2; + #pragma unroll 1 + for (int local_k_2 = 0; local_k_2 < 1; local_k_2++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 71304336;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0060.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0060.cu new file mode 100644 index 00000000..254d8e9d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0060.cu @@ -0,0 +1,81 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_highd_splitk_reduce_blockn64_g1r4_streamdep_r63_v1(float* __restrict__ partial_scores, int* __restrict__ partial_indices, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_slices) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid / 4; + int row_lane = tid % 4; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 64 + row; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_slice = row_lane; iter_slice < K_slices; iter_slice += 4) { + int partial_offset = (point_tile_idx * (unsigned int)K_slices + (unsigned int)iter_slice) * 64 + (unsigned int)row; + float score = partial_scores[partial_offset]; + int idx = partial_indices[partial_offset]; + if (score > best_score) { + best_score = score; + best_idx = idx; + } + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best_score, 1); + float peer_score = _shfl_xor_0; + int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, best_idx, 1); + int peer_idx = _shfl_xor_1; + if (peer_score > best_score) { + best_score = peer_score; + best_idx = peer_idx; + } + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, best_score, 2); + peer_score = _shfl_xor_2; + int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, best_idx, 2); + peer_idx = _shfl_xor_3; + if (peer_score > best_score) { + best_score = peer_score; + best_idx = peer_idx; + } + if (row_lane == 0) { + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = best_idx; + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0061.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0061.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0061.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0062.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0062.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0062.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0063.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0063.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0063.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0064.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0064.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0064.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0065.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0065.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0065.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0066.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0066.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0066.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0067.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0067.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0067.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0068.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0068.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0068.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0069.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0069.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0069.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0070.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0070.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0070.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0071.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0071.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0071.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0072.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0072.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0072.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0073.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0073.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0073.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0074.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0074.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0074.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0075.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0075.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0075.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0076.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0076.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0076.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0077.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0077.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0077.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0078.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0078.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0078.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0079.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0079.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0079.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0080.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0080.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0080.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0081.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0081.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0081.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0082.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0082.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0082.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0083.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0083.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0083.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0084.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0084.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0084.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0085.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0085.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0085.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0086.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0086.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0086.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0087.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0087.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0087.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0088.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0088.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0088.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0089.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0089.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0089.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0090.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0090.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0090.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0091.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0091.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0091.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0092.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0092.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0092.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0093.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0093.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0093.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0094.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0094.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0094.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0095.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0095.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0095.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0096.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0096.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0096.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0097.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0097.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0097.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0098.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0098.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0098.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0099.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0099.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0099.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0100.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0100.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0100.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0101.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0101.cu new file mode 100644 index 00000000..ea986bdf --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0101.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 24 * 8; + int row = vec_idx / 24; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 24 * 8; + int row_1 = vec_idx_1 / 24; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0102.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0102.cu new file mode 100644 index 00000000..7070fa90 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0102.cu @@ -0,0 +1,600 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 49152 +#define SMEM_SMEM_X_STRIDE 49152 +#define SMEM_SMEM_C_OFF 50176 +#define SMEM_SMEM_C_STAGE_BYTES 98304 +#define SMEM_SMEM_C_STRIDE 98304 +#define SMEM_SMEM_CSQ_OFF 148480 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 149504 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int smem_c_addr = smem + 50176; + float* smem_csq = reinterpret_cast(smem_raw + 148480); + const int smem_csq_addr = smem + 148480; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 49152); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 98304); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0103.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0103.cu new file mode 100644 index 00000000..ea986bdf --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0103.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 24 * 8; + int row = vec_idx / 24; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 24 * 8; + int row_1 = vec_idx_1 / 24; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0104.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0104.cu new file mode 100644 index 00000000..7070fa90 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0104.cu @@ -0,0 +1,600 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 49152 +#define SMEM_SMEM_X_STRIDE 49152 +#define SMEM_SMEM_C_OFF 50176 +#define SMEM_SMEM_C_STAGE_BYTES 98304 +#define SMEM_SMEM_C_STRIDE 98304 +#define SMEM_SMEM_CSQ_OFF 148480 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 149504 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int smem_c_addr = smem + 50176; + float* smem_csq = reinterpret_cast(smem_raw + 148480); + const int smem_csq_addr = smem + 148480; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 49152); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 98304); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0105.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0105.cu new file mode 100644 index 00000000..7070fa90 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0105.cu @@ -0,0 +1,600 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 49152 +#define SMEM_SMEM_X_STRIDE 49152 +#define SMEM_SMEM_C_OFF 50176 +#define SMEM_SMEM_C_STAGE_BYTES 98304 +#define SMEM_SMEM_C_STRIDE 98304 +#define SMEM_SMEM_CSQ_OFF 148480 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 149504 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int smem_c_addr = smem + 50176; + float* smem_csq = reinterpret_cast(smem_raw + 148480); + const int smem_csq_addr = smem + 148480; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 49152); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 98304); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0106.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0106.cu new file mode 100644 index 00000000..7070fa90 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0106.cu @@ -0,0 +1,600 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 49152 +#define SMEM_SMEM_X_STRIDE 49152 +#define SMEM_SMEM_C_OFF 50176 +#define SMEM_SMEM_C_STAGE_BYTES 98304 +#define SMEM_SMEM_C_STRIDE 98304 +#define SMEM_SMEM_CSQ_OFF 148480 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 149504 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int smem_c_addr = smem + 50176; + float* smem_csq = reinterpret_cast(smem_raw + 148480); + const int smem_csq_addr = smem + 148480; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 49152); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 98304); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0107.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0107.cu new file mode 100644 index 00000000..17caae28 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0107.cu @@ -0,0 +1,1211 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 49152 +#define SMEM_SMEM_X0_STRIDE 49152 +#define SMEM_SMEM_X1_OFF 50176 +#define SMEM_SMEM_X1_STAGE_BYTES 49152 +#define SMEM_SMEM_X1_STRIDE 49152 +#define SMEM_SMEM_C_OFF 99328 +#define SMEM_SMEM_C_STAGE_BYTES 98304 +#define SMEM_SMEM_C_STRIDE 98304 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int smem_x1_addr = smem + 50176; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); + const int smem_c_addr = smem + 99328; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 49152); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 49152); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 98304); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x0_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_4 = smem_x1_addr + 16384; + int _mma_a_lo_4 = make_warp_uniform((_mma_a_addr_4 >> 4) & 0x3FFF); + int _mma_b_addr_4 = smem_c_addr + 32768; + int _mma_b_lo_4 = make_warp_uniform((_mma_b_addr_4 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_5 = smem_x1_addr + 32768; + int _mma_a_lo_5 = make_warp_uniform((_mma_a_addr_5 >> 4) & 0x3FFF); + int _mma_b_addr_5 = smem_c_addr + 65536; + int _mma_b_lo_5 = make_warp_uniform((_mma_b_addr_5 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0108.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0108.cu new file mode 100644 index 00000000..17caae28 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0108.cu @@ -0,0 +1,1211 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 49152 +#define SMEM_SMEM_X0_STRIDE 49152 +#define SMEM_SMEM_X1_OFF 50176 +#define SMEM_SMEM_X1_STAGE_BYTES 49152 +#define SMEM_SMEM_X1_STRIDE 49152 +#define SMEM_SMEM_C_OFF 99328 +#define SMEM_SMEM_C_STAGE_BYTES 98304 +#define SMEM_SMEM_C_STRIDE 98304 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int smem_x1_addr = smem + 50176; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); + const int smem_c_addr = smem + 99328; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 49152); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 49152); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 98304); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x0_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_4 = smem_x1_addr + 16384; + int _mma_a_lo_4 = make_warp_uniform((_mma_a_addr_4 >> 4) & 0x3FFF); + int _mma_b_addr_4 = smem_c_addr + 32768; + int _mma_b_lo_4 = make_warp_uniform((_mma_b_addr_4 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_5 = smem_x1_addr + 32768; + int _mma_a_lo_5 = make_warp_uniform((_mma_a_addr_5 >> 4) & 0x3FFF); + int _mma_b_addr_5 = smem_c_addr + 65536; + int _mma_b_lo_5 = make_warp_uniform((_mma_b_addr_5 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0109.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0109.cu new file mode 100644 index 00000000..bf486f4b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0109.cu @@ -0,0 +1,639 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 65536 +#define SMEM_SMEM_X_STRIDE 65536 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 131072 +#define SMEM_SMEM_C_STRIDE 131072 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 65536); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 131072); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x_addr + 49152; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 98304; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0110.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0110.cu new file mode 100644 index 00000000..bf486f4b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0110.cu @@ -0,0 +1,639 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 65536 +#define SMEM_SMEM_X_STRIDE 65536 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 131072 +#define SMEM_SMEM_C_STRIDE 131072 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 65536); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 131072); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x_addr + 49152; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 98304; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0111.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0111.cu new file mode 100644 index 00000000..bf486f4b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0111.cu @@ -0,0 +1,639 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 65536 +#define SMEM_SMEM_X_STRIDE 65536 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 131072 +#define SMEM_SMEM_C_STRIDE 131072 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 65536); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 131072); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x_addr + 49152; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 98304; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0112.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0112.cu new file mode 100644 index 00000000..bf486f4b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0112.cu @@ -0,0 +1,639 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 65536 +#define SMEM_SMEM_X_STRIDE 65536 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 131072 +#define SMEM_SMEM_C_STRIDE 131072 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 65536); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 131072); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x_addr + 49152; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 98304; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0113.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0113.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0113.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0114.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0114.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0114.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0115.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0115.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0115.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0116.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0116.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0116.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0117.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0117.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0117.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0118.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0118.cu new file mode 100644 index 00000000..885f8507 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0118.cu @@ -0,0 +1,649 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM0_OFFSET 0 +#define TMEM_SCORE_TMEM1_OFFSET 256 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C0_OFF 9216 +#define SMEM_SMEM_C0_STAGE_BYTES 32768 +#define SMEM_SMEM_C0_STRIDE 32768 +#define SMEM_SMEM_C1_OFF 41984 +#define SMEM_SMEM_C1_STAGE_BYTES 32768 +#define SMEM_SMEM_C1_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 74752 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 75776 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_paired_xreuse_dualtmem_producer_r47_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, unsigned long long* __restrict__ partial_keys, int B, int N, int D, int K, int num_n_tiles, int K_tiles, int K_slices) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c0_addr = smem + 9216; + __nv_bfloat16* smem_c1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 41984); + const int smem_c1_addr = smem + 41984; + float* smem_csq = reinterpret_cast(smem_raw + 74752); + const int smem_csq_addr = smem + 74752; + + // Mbarrier init (8 groups, 8 barriers) + // Mbarriers at smem_raw[0..64) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 64); + if (warp == 0) { + int _tmem_hold = smem + 64; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c0_full_addr (mbar_base + 16) + #define c0_empty_addr (mbar_base + 24) + #define c1_full_addr (mbar_base + 32) + #define c1_empty_addr (mbar_base + 40) + #define score_full_addr (mbar_base + 48) + #define score_empty_addr (mbar_base + 56) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem0 = taddr; + const int tmem_score_tmem1 = taddr + 256; + + // ---- Role: compute ---- + if (warp <= 3) { + { // compute_main + int total_work = B * num_n_tiles * K_slices; + int compute_warp = warp; + int lane_pair = lane % 4; + int row_origin = compute_warp * 16; + int row_lane_base = row_origin + lane / 4; + int row0 = row_lane_base; + int row1 = row_lane_base + 8; + int compute_tid = compute_warp * 32 + lane; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_slice = work_idx % (unsigned int)K_slices; + int point_tile_idx = work_idx / (unsigned int)K_slices; + int batch = point_tile_idx / num_n_tiles; + int slice_k_start = iter_slice * 2; + int csq_smem_addr = smem_csq_addr; + float best0 = -3.4e+38f; + float best1 = -3.4e+38f; + int idx0 = slice_k_start * 256; + int idx1 = idx0; + int off_k0 = slice_k_start * 256; + int off_k1 = off_k0 + 256; + if ((float)compute_tid < 64.0f) { + int csq_base0 = compute_tid * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch * K + off_k0 + csq_base0); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half0[4]; + csq_half0[0] = 0.5f * _vec_load_0[0]; + csq_half0[1] = 0.5f * _vec_load_0[1]; + csq_half0[2] = 0.5f * _vec_load_0[2]; + csq_half0[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base0 * 4), "f"(csq_half0[0]), "f"(csq_half0[1]), "f"(csq_half0[2]), "f"(csq_half0[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + #pragma unroll + for (int score_base = 0; score_base < 256; score_base += 128) { + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, [%64];" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])), "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) + : "r"(taddr + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int rep = 0; rep < 16.0f; rep++) { + int local_reg = rep * 4; + int col_base = score_base + rep * 8 + lane_pair * 2; + float csq0 = smem_csq[col_base]; + float csq1 = smem_csq[col_base + 1]; + float d0 = _tmem_load_0[local_reg] - csq0; + if (d0 > best0) { + best0 = d0; + idx0 = off_k0 + col_base; + } + float d1 = _tmem_load_0[local_reg + 1] - csq1; + if (d1 > best0) { + best0 = d1; + idx0 = off_k0 + col_base + 1; + } + float d2 = _tmem_load_0[local_reg + 2] - csq0; + if (d2 > best1) { + best1 = d2; + idx1 = off_k0 + col_base; + } + float d3 = _tmem_load_0[local_reg + 3] - csq1; + if (d3 > best1) { + best1 = d3; + idx1 = off_k0 + col_base + 1; + } + } + } + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if ((float)compute_tid < 64.0f) { + int csq_base1 = compute_tid * 4; + float _vec_load_1[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch * K + off_k1 + csq_base1); + _vec_load_1[0 + 0] = _v4.x; + _vec_load_1[0 + 1] = _v4.y; + _vec_load_1[0 + 2] = _v4.z; + _vec_load_1[0 + 3] = _v4.w; + } + float csq_half1[4]; + csq_half1[0] = 0.5f * _vec_load_1[0]; + csq_half1[1] = 0.5f * _vec_load_1[1]; + csq_half1[2] = 0.5f * _vec_load_1[2]; + csq_half1[3] = 0.5f * _vec_load_1[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base1 * 4), "f"(csq_half1[0]), "f"(csq_half1[1]), "f"(csq_half1[2]), "f"(csq_half1[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + #pragma unroll + for (int score_base_1 = 0; score_base_1 < 256; score_base_1 += 128) { + float _tmem_load_1[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, [%64];" + : "=r"(*reinterpret_cast(&_tmem_load_1[0])), "=r"(*reinterpret_cast(&_tmem_load_1[1])), "=r"(*reinterpret_cast(&_tmem_load_1[2])), "=r"(*reinterpret_cast(&_tmem_load_1[3])), "=r"(*reinterpret_cast(&_tmem_load_1[4])), "=r"(*reinterpret_cast(&_tmem_load_1[5])), "=r"(*reinterpret_cast(&_tmem_load_1[6])), "=r"(*reinterpret_cast(&_tmem_load_1[7])), "=r"(*reinterpret_cast(&_tmem_load_1[8])), "=r"(*reinterpret_cast(&_tmem_load_1[9])), "=r"(*reinterpret_cast(&_tmem_load_1[10])), "=r"(*reinterpret_cast(&_tmem_load_1[11])), "=r"(*reinterpret_cast(&_tmem_load_1[12])), "=r"(*reinterpret_cast(&_tmem_load_1[13])), "=r"(*reinterpret_cast(&_tmem_load_1[14])), "=r"(*reinterpret_cast(&_tmem_load_1[15])), "=r"(*reinterpret_cast(&_tmem_load_1[16])), "=r"(*reinterpret_cast(&_tmem_load_1[17])), "=r"(*reinterpret_cast(&_tmem_load_1[18])), "=r"(*reinterpret_cast(&_tmem_load_1[19])), "=r"(*reinterpret_cast(&_tmem_load_1[20])), "=r"(*reinterpret_cast(&_tmem_load_1[21])), "=r"(*reinterpret_cast(&_tmem_load_1[22])), "=r"(*reinterpret_cast(&_tmem_load_1[23])), "=r"(*reinterpret_cast(&_tmem_load_1[24])), "=r"(*reinterpret_cast(&_tmem_load_1[25])), "=r"(*reinterpret_cast(&_tmem_load_1[26])), "=r"(*reinterpret_cast(&_tmem_load_1[27])), "=r"(*reinterpret_cast(&_tmem_load_1[28])), "=r"(*reinterpret_cast(&_tmem_load_1[29])), "=r"(*reinterpret_cast(&_tmem_load_1[30])), "=r"(*reinterpret_cast(&_tmem_load_1[31])), "=r"(*reinterpret_cast(&_tmem_load_1[32])), "=r"(*reinterpret_cast(&_tmem_load_1[33])), "=r"(*reinterpret_cast(&_tmem_load_1[34])), "=r"(*reinterpret_cast(&_tmem_load_1[35])), "=r"(*reinterpret_cast(&_tmem_load_1[36])), "=r"(*reinterpret_cast(&_tmem_load_1[37])), "=r"(*reinterpret_cast(&_tmem_load_1[38])), "=r"(*reinterpret_cast(&_tmem_load_1[39])), "=r"(*reinterpret_cast(&_tmem_load_1[40])), "=r"(*reinterpret_cast(&_tmem_load_1[41])), "=r"(*reinterpret_cast(&_tmem_load_1[42])), "=r"(*reinterpret_cast(&_tmem_load_1[43])), "=r"(*reinterpret_cast(&_tmem_load_1[44])), "=r"(*reinterpret_cast(&_tmem_load_1[45])), "=r"(*reinterpret_cast(&_tmem_load_1[46])), "=r"(*reinterpret_cast(&_tmem_load_1[47])), "=r"(*reinterpret_cast(&_tmem_load_1[48])), "=r"(*reinterpret_cast(&_tmem_load_1[49])), "=r"(*reinterpret_cast(&_tmem_load_1[50])), "=r"(*reinterpret_cast(&_tmem_load_1[51])), "=r"(*reinterpret_cast(&_tmem_load_1[52])), "=r"(*reinterpret_cast(&_tmem_load_1[53])), "=r"(*reinterpret_cast(&_tmem_load_1[54])), "=r"(*reinterpret_cast(&_tmem_load_1[55])), "=r"(*reinterpret_cast(&_tmem_load_1[56])), "=r"(*reinterpret_cast(&_tmem_load_1[57])), "=r"(*reinterpret_cast(&_tmem_load_1[58])), "=r"(*reinterpret_cast(&_tmem_load_1[59])), "=r"(*reinterpret_cast(&_tmem_load_1[60])), "=r"(*reinterpret_cast(&_tmem_load_1[61])), "=r"(*reinterpret_cast(&_tmem_load_1[62])), "=r"(*reinterpret_cast(&_tmem_load_1[63])) + : "r"(taddr + 256 + (unsigned int)score_base_1) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int rep_1 = 0; rep_1 < 16.0f; rep_1++) { + int local_reg_1 = rep_1 * 4; + int col_base_1 = score_base_1 + rep_1 * 8 + lane_pair * 2; + float csq0_1 = smem_csq[col_base_1]; + float csq1_1 = smem_csq[col_base_1 + 1]; + float d0_1 = _tmem_load_1[local_reg_1] - csq0_1; + if (d0_1 > best0) { + best0 = d0_1; + idx0 = off_k1 + col_base_1; + } + float d1_1 = _tmem_load_1[local_reg_1 + 1] - csq1_1; + if (d1_1 > best0) { + best0 = d1_1; + idx0 = off_k1 + col_base_1 + 1; + } + float d2_1 = _tmem_load_1[local_reg_1 + 2] - csq0_1; + if (d2_1 > best1) { + best1 = d2_1; + idx1 = off_k1 + col_base_1; + } + float d3_1 = _tmem_load_1[local_reg_1 + 3] - csq1_1; + if (d3_1 > best1) { + best1 = d3_1; + idx1 = off_k1 + col_base_1 + 1; + } + } + } + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best0, 1); + float peer0 = _shfl_xor_0; + int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, idx0, 1); + int peer0_idx = _shfl_xor_1; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer0_idx; + } + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, best0, 2); + peer0 = _shfl_xor_2; + int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, idx0, 2); + peer0_idx = _shfl_xor_3; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer0_idx; + } + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, best1, 1); + float peer1 = _shfl_xor_4; + int _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, idx1, 1); + int peer1_idx = _shfl_xor_5; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer1_idx; + } + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, best1, 2); + peer1 = _shfl_xor_6; + int _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, idx1, 2); + peer1_idx = _shfl_xor_7; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer1_idx; + } + if (lane_pair == 0) { + unsigned long long shift32 = 32; + unsigned long long mask64 = 4294967295; + uint32_t _amf_u_0 = __float_as_uint(best0); + uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; + unsigned int _amf_enc_0 = _amf_u_0 ^ _amf_mask_0; + unsigned long long key0 = (unsigned long long)_amf_enc_0 << shift32 | mask64 - (unsigned long long)idx0; + int partial_offset0 = work_idx * 64 + (unsigned int)row0; + *((unsigned long long*)(partial_keys + partial_offset0)) = key0; + uint32_t _amf_u_1 = __float_as_uint(best1); + uint32_t _amf_mask_1 = -int32_t(_amf_u_1 >> 31) | 0x80000000u; + unsigned int _amf_enc_1 = _amf_u_1 ^ _amf_mask_1; + unsigned long long key1 = (unsigned long long)_amf_enc_1 << shift32 | mask64 - (unsigned long long)idx1; + int partial_offset1 = work_idx * 64 + (unsigned int)row1; + *((unsigned long long*)(partial_keys + partial_offset1)) = key1; + } + } + } + // ---- Role: load ---- + } else if (warp == 4) { + { // load_main + int total_work_1 = B * num_n_tiles * K_slices; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c0_empty_0 = 1; + unsigned int _phase_c1_empty_0 = 1; + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_1; work_idx_1 += num_bids) { + int iter_slice_1 = work_idx_1 % (unsigned int)K_slices; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_slices; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int n_tile = point_tile_idx_1 % num_n_tiles; + int off_n = n_tile * 64; + int x_row = batch_1 * N + off_n; + int slice_k_start_1 = iter_slice_1 * 2; + int iter_k0 = slice_k_start_1; + int iter_k1 = slice_k_start_1 + 1; + int c_row0 = batch_1 * K + iter_k0 * 256; + int c_row1 = batch_1 * K + iter_k1 * 256; + #pragma unroll 7 + for (int feat_tile = 0; feat_tile < 7; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c0_empty_addr, _phase_c0_empty_0); + _phase_c0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c0_addr, c_tmap, 0, c_row0, feat_tile, c0_full_addr); + mbarrier_arrive_expect_tx(c0_full_addr, 32768); + mbarrier_wait(c1_empty_addr, _phase_c1_empty_0); + _phase_c1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c1_addr, c_tmap, 0, c_row1, feat_tile, c1_full_addr); + mbarrier_arrive_expect_tx(c1_full_addr, 32768); + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 5) { + { // mma_main + int total_work_2 = B * num_n_tiles * K_slices; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c0_full_0 = 0; + unsigned int _phase_c1_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = bid; work_idx_2 < total_work_2; work_idx_2 += num_bids) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 7 + for (int feat_tile_1 = 0; feat_tile_1 < 7; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c0_full_addr, _phase_c0_full_0); + _phase_c0_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c0_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 71304336;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem0), "r"(((init_flag) ? 0 : 1))); + elect_commit(c0_empty_addr); + mbarrier_wait(c1_full_addr, _phase_c1_full_0); + _phase_c1_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c1_addr; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 71304336;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem1), "r"(((init_flag) ? 0 : 1))); + elect_commit(c1_empty_addr); + elect_commit(x_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0119.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0119.cu new file mode 100644 index 00000000..c94c14c1 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0119.cu @@ -0,0 +1,63 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 64 + +#include + +extern "C" { + +__global__ __launch_bounds__(64) void +kernel_flash_kmeans_assign_highd_paired_ownerreduce_r39_reduce1_unroll_v1(unsigned long long* __restrict__ partial_keys, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_slices) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 64 + row; + unsigned long long best_key = 0; + #pragma unroll 8 + for (int iter_slice = 0; iter_slice < 8; iter_slice++) { + int partial_offset = (point_tile_idx * (unsigned int)K_slices + (unsigned int)iter_slice) * 64 + (unsigned int)row; + unsigned long long key = partial_keys[partial_offset]; + if (key > best_key) { + best_key = key; + } + } + unsigned long long mask64 = 4294967295; + unsigned long long inv_idx = best_key & mask64; + unsigned long long idx_u64 = mask64 - inv_idx; + int idx = (int)idx_u64; + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = idx; + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0120.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0120.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0120.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0121.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0121.cu new file mode 100644 index 00000000..0d005883 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0121.cu @@ -0,0 +1,543 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 41984 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 43008 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_paired_packedpartial_producer_7b3c_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, unsigned long long* __restrict__ partial_keys, int B, int N, int D, int K, int num_n_tiles, int K_tiles, int K_slices) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 41984); + const int smem_csq_addr = smem + 41984; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: compute ---- + if (warp <= 3) { + { // compute_main + int total_work = B * num_n_tiles * K_slices; + int compute_warp = warp; + int lane_pair = lane % 4; + int row_origin = compute_warp * 16; + int row_lane_base = row_origin + lane / 4; + int row0 = row_lane_base; + int row1 = row_lane_base + 8; + int compute_tid = compute_warp * 32 + lane; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_slice = work_idx % (unsigned int)K_slices; + int point_tile_idx = work_idx / (unsigned int)K_slices; + int batch = point_tile_idx / num_n_tiles; + int slice_k_start = iter_slice * 2; + int csq_smem_addr = smem_csq_addr; + float best0 = -3.4e+38f; + float best1 = -3.4e+38f; + int idx0 = slice_k_start * 256; + int idx1 = idx0; + #pragma unroll 1 + for (int local_k = 0; local_k < 2; local_k++) { + int iter_k = slice_k_start + local_k; + int off_k = iter_k * 256; + if ((float)compute_tid < 64.0f) { + int csq_base = compute_tid * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch * K + off_k + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + #pragma unroll + for (int score_base = 0; score_base < 256; score_base += 128) { + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, [%64];" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])), "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) + : "r"(taddr + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int rep = 0; rep < 16.0f; rep++) { + int local_reg = rep * 4; + int col_base = score_base + rep * 8 + lane_pair * 2; + float csq0 = smem_csq[col_base]; + float csq1 = smem_csq[col_base + 1]; + float d0 = _tmem_load_0[local_reg] - csq0; + if (d0 > best0) { + best0 = d0; + idx0 = off_k + col_base; + } + float d1 = _tmem_load_0[local_reg + 1] - csq1; + if (d1 > best0) { + best0 = d1; + idx0 = off_k + col_base + 1; + } + float d2 = _tmem_load_0[local_reg + 2] - csq0; + if (d2 > best1) { + best1 = d2; + idx1 = off_k + col_base; + } + float d3 = _tmem_load_0[local_reg + 3] - csq1; + if (d3 > best1) { + best1 = d3; + idx1 = off_k + col_base + 1; + } + } + } + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best0, 1); + float peer0 = _shfl_xor_0; + int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, idx0, 1); + int peer0_idx = _shfl_xor_1; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer0_idx; + } + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, best0, 2); + peer0 = _shfl_xor_2; + int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, idx0, 2); + peer0_idx = _shfl_xor_3; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer0_idx; + } + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, best1, 1); + float peer1 = _shfl_xor_4; + int _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, idx1, 1); + int peer1_idx = _shfl_xor_5; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer1_idx; + } + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, best1, 2); + peer1 = _shfl_xor_6; + int _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, idx1, 2); + peer1_idx = _shfl_xor_7; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer1_idx; + } + if (lane_pair == 0) { + unsigned long long shift32 = 32; + unsigned long long mask64 = 4294967295; + uint32_t _amf_u_0 = __float_as_uint(best0); + uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; + unsigned int _amf_enc_0 = _amf_u_0 ^ _amf_mask_0; + unsigned long long key0 = (unsigned long long)_amf_enc_0 << shift32 | mask64 - (unsigned long long)idx0; + int partial_offset0 = work_idx * 64 + (unsigned int)row0; + *((unsigned long long*)(partial_keys + partial_offset0)) = key0; + uint32_t _amf_u_1 = __float_as_uint(best1); + uint32_t _amf_mask_1 = -int32_t(_amf_u_1 >> 31) | 0x80000000u; + unsigned int _amf_enc_1 = _amf_u_1 ^ _amf_mask_1; + unsigned long long key1 = (unsigned long long)_amf_enc_1 << shift32 | mask64 - (unsigned long long)idx1; + int partial_offset1 = work_idx * 64 + (unsigned int)row1; + *((unsigned long long*)(partial_keys + partial_offset1)) = key1; + } + } + } + // ---- Role: load ---- + } else if (warp == 4) { + { // load_main + int total_work_1 = B * num_n_tiles * K_slices; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_1; work_idx_1 += num_bids) { + int iter_slice_1 = work_idx_1 % (unsigned int)K_slices; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_slices; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int n_tile = point_tile_idx_1 % num_n_tiles; + int off_n = n_tile * 64; + int x_row = batch_1 * N + off_n; + int slice_k_start_1 = iter_slice_1 * 2; + #pragma unroll 1 + for (int local_k_1 = 0; local_k_1 < 2; local_k_1++) { + int iter_k_1 = slice_k_start_1 + local_k_1; + int off_k_1 = iter_k_1 * 256; + int c_row = batch_1 * K + off_k_1; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 5) { + { // mma_main + int total_work_2 = B * num_n_tiles * K_slices; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = bid; work_idx_2 < total_work_2; work_idx_2 += num_bids) { + int iter_slice_2 = work_idx_2 % (unsigned int)K_slices; + int slice_k_start_2 = iter_slice_2 * 2; + #pragma unroll 1 + for (int local_k_2 = 0; local_k_2 < 2; local_k_2++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 71304336;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0122.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0122.cu new file mode 100644 index 00000000..cf89ea2d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0122.cu @@ -0,0 +1,71 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 128 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flash_kmeans_assign_highd_paired_packedpartial_reduce_r2_7b3c_v1(unsigned long long* __restrict__ partial_keys, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_slices) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid / 2; + int row_lane = tid % 2; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 64 + row; + unsigned long long best_key = 0; + #pragma unroll 1 + for (int iter_slice = row_lane; iter_slice < K_slices; iter_slice += 2) { + int partial_offset = (point_tile_idx * (unsigned int)K_slices + (unsigned int)iter_slice) * 64 + (unsigned int)row; + unsigned long long key = partial_keys[partial_offset]; + if (key > best_key) { + best_key = key; + } + } + unsigned long long _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best_key, 1); + unsigned long long peer_key = _shfl_xor_0; + if (peer_key > best_key) { + best_key = peer_key; + } + if (row_lane == 0) { + unsigned long long mask64 = 4294967295; + unsigned long long inv_idx = best_key & mask64; + unsigned long long idx_u64 = mask64 - inv_idx; + int idx = (int)idx_u64; + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = idx; + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0123.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0123.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0123.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0124.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0124.cu new file mode 100644 index 00000000..5b836885 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0124.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_microdim_pack_6cd2_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 8 * 8; + int row = vec_idx / 8; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 8 * 8; + int row_1 = vec_idx_1 / 8; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0125.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0125.cu new file mode 100644 index 00000000..1c18f40d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0125.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_microdim_6cd2_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0126.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0126.cu new file mode 100644 index 00000000..3e4e2222 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0126.cu @@ -0,0 +1,483 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORES_OFFSET 0 +#define NUM_X_PIPE_STAGES 3 +#define NUM_C_PIPE_STAGES 4 +#define NUM_SCORE_PIPE_STAGES 4 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 4096 +#define SMEM_SMEM_X_STRIDE 4096 +#define SMEM_SMEM_C_OFF 13312 +#define SMEM_SMEM_C_STAGE_BYTES 2048 +#define SMEM_SMEM_C_STRIDE 2048 +#define SMEM_SMEM_CSQ_OFF 21504 +#define SMEM_SMEM_CSQ_STAGE_BYTES 2048 +#define SMEM_SMEM_CSQ_STRIDE 2048 +#define SMEM_TOTAL 23552 +#define THREADS 192 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_microdim_d16_pipeline4_08f9_v4(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 13312); + const int smem_c_addr = smem + 13312; + float* smem_csq = reinterpret_cast(smem_raw + 21504); + const int smem_csq_addr = smem + 21504; + + // Mbarrier init (6 groups, 22 barriers) + // Mbarriers at smem_raw[0..176) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // --- pipeline 'x_pipe' --- + // x_full: 3 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + mbarrier_init_pred(smem + 8, 1, leader); + mbarrier_init_pred(smem + 16, 1, leader); + // x_empty: 3 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + mbarrier_init_pred(smem + 32, 1, leader); + mbarrier_init_pred(smem + 40, 1, leader); + // --- pipeline 'c_pipe' --- + // c_full: 4 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + mbarrier_init_pred(smem + 56, 1, leader); + mbarrier_init_pred(smem + 64, 1, leader); + mbarrier_init_pred(smem + 72, 1, leader); + // c_empty: 4 barriers, init_count=1 + mbarrier_init_pred(smem + 80, 1, leader); + mbarrier_init_pred(smem + 88, 1, leader); + mbarrier_init_pred(smem + 96, 1, leader); + mbarrier_init_pred(smem + 104, 1, leader); + // --- pipeline 'score_pipe' --- + // score_full: 4 barriers, init_count=1 + mbarrier_init_pred(smem + 112, 1, leader); + mbarrier_init_pred(smem + 120, 1, leader); + mbarrier_init_pred(smem + 128, 1, leader); + mbarrier_init_pred(smem + 136, 1, leader); + // score_empty: 4 barriers, init_count=4 + mbarrier_init_pred(smem + 144, 4, leader); + mbarrier_init_pred(smem + 152, 4, leader); + mbarrier_init_pred(smem + 160, 4, leader); + mbarrier_init_pred(smem + 168, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 176); + if (warp == 0) { + int _tmem_hold = smem + 176; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 48) + #define c_empty_addr (mbar_base + 80) + #define score_full_addr (mbar_base + 112) + #define score_empty_addr (mbar_base + 144) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_scores = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + unsigned int x_stage = 0; + unsigned int c_stage = 0; + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty = 1; + unsigned int _phase_c_empty = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int x_row = batch * N + n_tile * 128; + mbarrier_wait(x_empty_addr + (x_stage) * 8, _phase_x_empty); + tma_3d_gmem2smem(smem_x_addr + x_stage * 4096, x_tmap, 0, x_row, 0, x_full_addr + (x_stage) * 8); + mbarrier_arrive_expect_tx(x_full_addr + (x_stage) * 8, 4096); + x_stage = (x_stage + 1) % 3; + if (x_stage == 0) { _phase_x_empty ^= 1; } + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int c_row = batch * K + iter_k * 64; + mbarrier_wait(c_empty_addr + (c_stage) * 8, _phase_c_empty); + tma_3d_gmem2smem(smem_c_addr + c_stage * 2048, c_tmap, 0, c_row, 0, c_full_addr + (c_stage) * 8); + mbarrier_arrive_expect_tx(c_full_addr + (c_stage) * 8, 2048); + c_stage = (c_stage + 1) % 4; + if (c_stage == 0) { _phase_c_empty ^= 1; } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + unsigned int x_stage_1 = 0; + unsigned int c_stage_1 = 0; + unsigned int score_stage = 0; + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full = 0; + unsigned int _phase_c_full = 0; + unsigned int _phase_score_empty = 1; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr + (x_stage_1) * 8, _phase_x_full); + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(c_full_addr + (c_stage_1) * 8, _phase_c_full); + mbarrier_wait(score_empty_addr + (score_stage) * 8, _phase_score_empty); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr + x_stage_1 * 4096; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr + c_stage_1 * 2048; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + mma_ss_step(_mma_a_lo_0, _mma_b_lo_0, (tmem_scores + (score_stage * 64)), 135267472, 0, 0xC0004010U, 0xC0004010U); + elect_commit(score_full_addr + (score_stage) * 8); + elect_commit(c_empty_addr + (c_stage_1) * 8); + c_stage_1 = (c_stage_1 + 1) % 4; + if (c_stage_1 == 0) { _phase_c_full ^= 1; } + score_stage = (score_stage + 1) % 4; + if (score_stage == 0) { _phase_score_empty ^= 1; } + } + if (elect_sync()) { + mbarrier_arrive(x_empty_addr + (x_stage_1) * 8); + } + x_stage_1 = (x_stage_1 + 1) % 3; + if (x_stage_1 == 0) { _phase_x_full ^= 1; } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + unsigned int score_stage_1 = 0; + int num_tiles_2 = B * num_n_tiles; + int row = warp % 4 * 32 + lane; + int row_base = warp % 4 * 32 << 16; + unsigned int _phase_score_full = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int out_offset = batch_1 * N + n_tile_1 * 128 + row; + float best0 = -3.4e+38f; + float best1 = -3.4e+38f; + float best2 = -3.4e+38f; + float best3 = -3.4e+38f; + int idx0 = 0; + int idx1 = 0; + int idx2 = 0; + int idx3 = 0; + int csq_base = row * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(smem_csq_addr + (unsigned int)(csq_base * 4)), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int k_base = iter_k_1 * 64; + mbarrier_wait(score_full_addr + (score_stage_1) * 8, _phase_score_full); + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x64.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, [%64];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]) + : "r"(taddr + (unsigned int)row_base + score_stage_1 * 64) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int j_base = 0; j_base < 64; j_base += 4) { + float d0 = _tmem_load_0[j_base] - smem_csq[k_base + j_base]; + if (d0 > best0) { + best0 = d0; + idx0 = k_base + j_base; + } + float d1 = _tmem_load_0[j_base + 1] - smem_csq[k_base + j_base + 1]; + if (d1 > best1) { + best1 = d1; + idx1 = k_base + j_base + 1; + } + float d2 = _tmem_load_0[j_base + 2] - smem_csq[k_base + j_base + 2]; + if (d2 > best2) { + best2 = d2; + idx2 = k_base + j_base + 2; + } + float d3 = _tmem_load_0[j_base + 3] - smem_csq[k_base + j_base + 3]; + if (d3 > best3) { + best3 = d3; + idx3 = k_base + j_base + 3; + } + } + if (elect_sync()) { + mbarrier_arrive(score_empty_addr + (score_stage_1) * 8); + } + score_stage_1 = (score_stage_1 + 1) % 4; + if (score_stage_1 == 0) { _phase_score_full ^= 1; } + } + if (best1 > best0) { + best0 = best1; + idx0 = idx1; + } + if (best3 > best2) { + best2 = best3; + idx2 = idx3; + } + if (best2 > best0) { + best0 = best2; + idx0 = idx2; + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + *((int*)(out + out_offset)) = idx0; + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0127.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0127.cu new file mode 100644 index 00000000..5b836885 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0127.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_microdim_pack_6cd2_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 8 * 8; + int row = vec_idx / 8; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 8 * 8; + int row_1 = vec_idx_1 / 8; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0128.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0128.cu new file mode 100644 index 00000000..1c18f40d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0128.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_microdim_6cd2_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0129.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0129.cu new file mode 100644 index 00000000..3778419b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0129.cu @@ -0,0 +1,520 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 2048 +#define SMEM_SMEM_CSQ_STRIDE 2048 +#define SMEM_TOTAL 52224 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_microdim_raw_tma_08f9_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_2d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_2d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (off_k_1 + score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (off_k_1 + score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0130.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0130.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0130.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0131.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0131.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0131.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0132.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0132.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0132.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0133.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0133.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0133.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0134.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0134.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0134.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0135.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0135.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0135.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0136.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0136.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0136.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0137.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0137.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0137.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0138.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0138.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0138.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0139.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0139.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0139.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0140.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0140.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0140.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0141.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0141.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0141.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0142.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0142.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0142.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0143.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0143.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0143.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0144.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0144.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0144.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0145.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0145.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0145.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0146.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0146.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0146.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0147.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0147.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0147.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0148.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0148.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0148.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0149.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0149.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0149.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0150.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0150.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0150.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0151.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0151.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0151.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0152.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0152.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0152.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0153.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0153.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0153.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0154.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0154.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0154.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0155.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0155.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0155.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0156.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0156.cu new file mode 100644 index 00000000..7ed658bb --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0156.cu @@ -0,0 +1,562 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_X_RAW_OFF 1024 +#define SMEM_X_RAW_STAGE_BYTES 14336 +#define SMEM_X_RAW_STRIDE 14336 +#define SMEM_C_DIRECT00_OFF 15360 +#define SMEM_C_DIRECT00_STAGE_BYTES 1792 +#define SMEM_C_DIRECT00_STRIDE 1792 +#define SMEM_C_DIRECT01_OFF 17152 +#define SMEM_C_DIRECT01_STAGE_BYTES 1792 +#define SMEM_C_DIRECT01_STRIDE 1792 +#define SMEM_C_DIRECT10_OFF 18944 +#define SMEM_C_DIRECT10_STAGE_BYTES 1792 +#define SMEM_C_DIRECT10_STRIDE 1792 +#define SMEM_C_DIRECT11_OFF 20736 +#define SMEM_C_DIRECT11_STAGE_BYTES 1792 +#define SMEM_C_DIRECT11_STRIDE 1792 +#define SMEM_SX_OFF 22528 +#define SMEM_SX_STAGE_BYTES 14336 +#define SMEM_SX_STRIDE 14336 +#define SMEM_SS_OFF 36864 +#define SMEM_SS_STAGE_BYTES 4096 +#define SMEM_SS_STRIDE 4096 +#define SMEM_GROUP_KEYS_OFF 40960 +#define SMEM_GROUP_KEYS_STAGE_BYTES 1024 +#define SMEM_GROUP_KEYS_STRIDE 1024 +#define SMEM_LOCAL_KEYS_OFF 41984 +#define SMEM_LOCAL_KEYS_STAGE_BYTES 512 +#define SMEM_LOCAL_KEYS_STRIDE 512 +#define SMEM_CLUSTER_KEYS_OFF 42496 +#define SMEM_CLUSTER_KEYS_STAGE_BYTES 4096 +#define SMEM_CLUSTER_KEYS_STRIDE 4096 +#define SMEM_TOTAL 46592 +#define THREADS 256 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ void cp_async_bulk_gmem2smem( + unsigned smem_addr, const void* gmem_ptr, unsigned bytes, int mbar_addr) { + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(smem_addr), "l"(gmem_ptr), "r"(bytes), "r"(mbar_addr) + : "memory"); +} + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 8; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 8; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* x_raw = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int x_raw_addr = smem + 1024; + __nv_bfloat16* c_direct00 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 15360); + const int c_direct00_addr = smem + 15360; + __nv_bfloat16* c_direct01 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17152); + const int c_direct01_addr = smem + 17152; + __nv_bfloat16* c_direct10 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 18944); + const int c_direct10_addr = smem + 18944; + __nv_bfloat16* c_direct11 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 20736); + const int c_direct11_addr = smem + 20736; + __nv_bfloat16* sx = reinterpret_cast<__nv_bfloat16*>(smem_raw + 22528); + const int sx_addr = smem + 22528; + float* ss = reinterpret_cast(smem_raw + 36864); + const int ss_addr = smem + 36864; + unsigned long long* group_keys = reinterpret_cast(smem_raw + 40960); + const int group_keys_addr = smem + 40960; + unsigned long long* local_keys = reinterpret_cast(smem_raw + 41984); + const int local_keys_addr = smem + 41984; + unsigned long long* cluster_keys = reinterpret_cast(smem_raw + 42496); + const int cluster_keys_addr = smem + 42496; + + // Mbarrier init (6 groups, 13 barriers) + // Mbarriers at smem_raw[0..104) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_ready: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // c_ready00: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_ready01: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_ready10: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_ready11: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // keys_ready: 8 barriers, init_count=8 + mbarrier_init_pred(smem + 40, 8, leader); + mbarrier_init_pred(smem + 48, 8, leader); + mbarrier_init_pred(smem + 56, 8, leader); + mbarrier_init_pred(smem + 64, 8, leader); + mbarrier_init_pred(smem + 72, 8, leader); + mbarrier_init_pred(smem + 80, 8, leader); + mbarrier_init_pred(smem + 88, 8, leader); + mbarrier_init_pred(smem + 96, 8, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); + + const int mbar_base = smem; + #define x_ready_addr (mbar_base + 0) + #define c_ready00_addr (mbar_base + 8) + #define c_ready01_addr (mbar_base + 16) + #define c_ready10_addr (mbar_base + 24) + #define c_ready11_addr (mbar_base + 32) + #define keys_ready_addr (mbar_base + 40) + + // === Task calls (dependency order) === + int total_tiles = B * num_n_tiles; + unsigned int _phase_x_ready_0 = 0; + unsigned int _phase_c_ready00_0 = 0; + unsigned int _phase_c_ready01_0 = 0; + unsigned int _phase_c_ready10_0 = 0; + unsigned int _phase_c_ready11_0 = 0; + unsigned int _phase_keys_ready_0 = 0; + #pragma unroll 1 + for (unsigned int tile = cluster_id; tile < total_tiles; tile += num_clusters) { + int batch = tile / (unsigned int)num_n_tiles; + int nt = tile % (unsigned int)num_n_tiles; + int point_base = batch * N + nt * 64; + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(x_ready_addr, 14336); + cp_async_bulk_gmem2smem(x_raw_addr, reinterpret_cast(reinterpret_cast(x) + ((unsigned long long)(point_base * D) * (unsigned long long)2)), 14336, x_ready_addr); + } + } + mbarrier_wait(x_ready_addr, _phase_x_ready_0); + _phase_x_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll 1 + for (unsigned int i = tid; i < 7168; i += 256) { + int row = i / 112; + int col = i % 112; + { + __nv_bfloat16 _bval_3404531696 = __float2bfloat16_rn(x_raw[i]); + uint16_t _bits_3404531696 = *(uint16_t*)&_bval_3404531696; + uint32_t _addr_3404531696 = static_cast((sx_addr + (unsigned int)(row * 224 + col * 2))); + asm volatile("st.shared.b16 [%0], %1;" :: "r"(_addr_3404531696), "h"(_bits_3404531696) : "memory"); + } + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + int warp_id_in_role = (warp - 0); + int group = warp_id_in_role / 4; + int local_warp = warp_id_in_role % 4; + int local_row = lane % 16; + int row_base = local_warp * 16; + float best = -3.4e+38f; + int owner_k_tiles = K / 64; + int group_k_tiles = owner_k_tiles / 2; + int group_tile_begin = cta_rank * owner_k_tiles + group * group_k_tiles; + int best_idx = group_tile_begin * 8; + int first_kbase = group_tile_begin * 8; + if (group == 0) { + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready00_addr, 1792); + cp_async_bulk_gmem2smem(c_direct00_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + first_kbase) * D) * (unsigned long long)2)), 1792, c_ready00_addr); + } + } + } else if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready10_addr, 1792); + cp_async_bulk_gmem2smem(c_direct10_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + first_kbase) * D) * (unsigned long long)2)), 1792, c_ready10_addr); + } + } + #pragma unroll 1 + for (int group_kt = 0; group_kt < group_k_tiles; group_kt++) { + int current_stage = group_kt % 2; + int kt = group_tile_begin + group_kt; + int kbase = kt * 8; + bool has_next = group_k_tiles > group_kt + 1; + if (has_next) { + int next_stage = (group_kt + 1) % 2; + int next_kbase = (kt + 1) * 8; + if (group == 0) { + if (next_stage == 0) { + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready00_addr, 1792); + cp_async_bulk_gmem2smem(c_direct00_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready00_addr); + } + } + } else if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready01_addr, 1792); + cp_async_bulk_gmem2smem(c_direct01_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready01_addr); + } + } + } else if (next_stage == 0) { + if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready10_addr, 1792); + cp_async_bulk_gmem2smem(c_direct10_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready10_addr); + } + } + } else { + if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready11_addr, 1792); + cp_async_bulk_gmem2smem(c_direct11_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready11_addr); + } + } + } + } + if (group == 0) { + if (current_stage == 0) { + mbarrier_wait(c_ready00_addr, _phase_c_ready00_0); + _phase_c_ready00_0 ^= 1; + } else { + mbarrier_wait(c_ready01_addr, _phase_c_ready01_0); + _phase_c_ready01_0 ^= 1; + } + } else if (current_stage == 0) { + mbarrier_wait(c_ready10_addr, _phase_c_ready10_0); + _phase_c_ready10_0 ^= 1; + } else { + mbarrier_wait(c_ready11_addr, _phase_c_ready11_0); + _phase_c_ready11_0 ^= 1; + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + unsigned int a[4]; + unsigned int b[2]; + float acc[4]; + unsigned int a_addr = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + lane / 16 * 16)); + unsigned int b_addr = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[0]), "=f"(acc[1]), "=f"(acc[2]), "=f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_0 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 32))); + unsigned int b_addr_1 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_0) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_1) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_2 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 64))); + unsigned int b_addr_3 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_2) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_3) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_4 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 96))); + unsigned int b_addr_5 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_4) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_5) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_6 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 128))); + unsigned int b_addr_7 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_6) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_7) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_8 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 160))); + unsigned int b_addr_9 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_8) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_9) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_10 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 192))); + unsigned int b_addr_11 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_10) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_11) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + #pragma unroll + for (int rp = 0; rp < 2; rp++) { + #pragma unroll + for (int cp = 0; cp < 2; cp++) { + int rr = lane / 4 + rp * 8; + int cc = lane % 4 * 2 + cp; + { + uint32_t _addr_3405188304 = static_cast((ss_addr + (unsigned int)((group * 64 + row_base + rr) * 32 + cc * 4))); + asm volatile("st.shared.f32 [%0], %1;" :: "r"(_addr_3405188304), "f"(acc[rp * 2 + cp]) : "memory"); + } + } + } + __syncwarp(); + if (lane < 16) { + #pragma unroll + for (int kk = 0; kk < 8; kk++) { + float score = ss[(group * 64 + row_base + local_row) * 8 + kk] - 0.5f * c_sq[batch * K + kbase + kk]; + if (score > best) { + best = score; + best_idx = kbase + kk; + } + } + } + if (group == 0) { + asm volatile("barrier.sync 1, %0;" :: "r"(128)); + } else { + asm volatile("barrier.sync 2, %0;" :: "r"(128)); + } + } + if (lane < 16) { + uint32_t _amf_u_0 = __float_as_uint(best); + uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; + unsigned int _amf_enc_0 = _amf_u_0 ^ _amf_mask_0; + unsigned long long shift64 = 32; + unsigned long long mask64 = 4294967295; + group_keys[group * 64 + row_base + local_row] = (unsigned long long)_amf_enc_0 << shift64 | mask64 - (unsigned long long)best_idx; + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + if (group == 0 && lane < 16) { + unsigned long long first_key = group_keys[row_base + local_row]; + unsigned long long second_key = group_keys[64 + row_base + local_row]; + local_keys[row_base + local_row] = ((second_key > first_key) ? second_key : first_key); + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + if (warp == 0) { + if (elect_sync()) { + uint32_t _mapa_0; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_0) : "r"(cluster_keys_addr + (unsigned int)(cta_rank * 512)), "r"(0)); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + uint32_t _mapa_1; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_1) : "r"(keys_ready_addr), "r"(0)); + asm volatile( + "mbarrier.arrive.expect_tx.release.cluster.shared::cluster.b64 _, [%0], %1;" + :: "r"(_mapa_1), "r"((uint32_t)(512)) : "memory"); + uint32_t _mapa_2; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_2) : "r"(keys_ready_addr), "r"(0)); + asm volatile( + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(_mapa_0), "r"(local_keys_addr), "r"((uint32_t)(512)), "r"(_mapa_2) + : "memory"); + } + } + if (cta_rank == 0) { + mbarrier_wait(keys_ready_addr, _phase_keys_ready_0); + _phase_keys_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int row_1 = tid / 2; + int lane_pair = tid % 2; + if (row_1 < 64) { + unsigned long long best_key = 0; + #pragma unroll + for (int peer = lane_pair; peer < 8; peer += 2) { + unsigned long long key = cluster_keys[peer * 64 + row_1]; + if (key > best_key) { + best_key = key; + } + } + unsigned long long _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best_key, 1); + unsigned long long peer_key = _shfl_xor_0; + if (peer_key > best_key) { + best_key = peer_key; + } + if (lane_pair == 0) { + unsigned long long mask64_1 = 4294967295; + int idx = (int)(mask64_1 - (best_key & mask64_1)); + *((int*)(out + (batch * N + nt * 64 + row_1))) = idx; + } + } + } + } + + // Cleanup + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0157.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0157.cu new file mode 100644 index 00000000..7ed658bb --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0157.cu @@ -0,0 +1,562 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_X_RAW_OFF 1024 +#define SMEM_X_RAW_STAGE_BYTES 14336 +#define SMEM_X_RAW_STRIDE 14336 +#define SMEM_C_DIRECT00_OFF 15360 +#define SMEM_C_DIRECT00_STAGE_BYTES 1792 +#define SMEM_C_DIRECT00_STRIDE 1792 +#define SMEM_C_DIRECT01_OFF 17152 +#define SMEM_C_DIRECT01_STAGE_BYTES 1792 +#define SMEM_C_DIRECT01_STRIDE 1792 +#define SMEM_C_DIRECT10_OFF 18944 +#define SMEM_C_DIRECT10_STAGE_BYTES 1792 +#define SMEM_C_DIRECT10_STRIDE 1792 +#define SMEM_C_DIRECT11_OFF 20736 +#define SMEM_C_DIRECT11_STAGE_BYTES 1792 +#define SMEM_C_DIRECT11_STRIDE 1792 +#define SMEM_SX_OFF 22528 +#define SMEM_SX_STAGE_BYTES 14336 +#define SMEM_SX_STRIDE 14336 +#define SMEM_SS_OFF 36864 +#define SMEM_SS_STAGE_BYTES 4096 +#define SMEM_SS_STRIDE 4096 +#define SMEM_GROUP_KEYS_OFF 40960 +#define SMEM_GROUP_KEYS_STAGE_BYTES 1024 +#define SMEM_GROUP_KEYS_STRIDE 1024 +#define SMEM_LOCAL_KEYS_OFF 41984 +#define SMEM_LOCAL_KEYS_STAGE_BYTES 512 +#define SMEM_LOCAL_KEYS_STRIDE 512 +#define SMEM_CLUSTER_KEYS_OFF 42496 +#define SMEM_CLUSTER_KEYS_STAGE_BYTES 4096 +#define SMEM_CLUSTER_KEYS_STRIDE 4096 +#define SMEM_TOTAL 46592 +#define THREADS 256 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ void cp_async_bulk_gmem2smem( + unsigned smem_addr, const void* gmem_ptr, unsigned bytes, int mbar_addr) { + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(smem_addr), "l"(gmem_ptr), "r"(bytes), "r"(mbar_addr) + : "memory"); +} + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 8; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 8; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* x_raw = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int x_raw_addr = smem + 1024; + __nv_bfloat16* c_direct00 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 15360); + const int c_direct00_addr = smem + 15360; + __nv_bfloat16* c_direct01 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17152); + const int c_direct01_addr = smem + 17152; + __nv_bfloat16* c_direct10 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 18944); + const int c_direct10_addr = smem + 18944; + __nv_bfloat16* c_direct11 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 20736); + const int c_direct11_addr = smem + 20736; + __nv_bfloat16* sx = reinterpret_cast<__nv_bfloat16*>(smem_raw + 22528); + const int sx_addr = smem + 22528; + float* ss = reinterpret_cast(smem_raw + 36864); + const int ss_addr = smem + 36864; + unsigned long long* group_keys = reinterpret_cast(smem_raw + 40960); + const int group_keys_addr = smem + 40960; + unsigned long long* local_keys = reinterpret_cast(smem_raw + 41984); + const int local_keys_addr = smem + 41984; + unsigned long long* cluster_keys = reinterpret_cast(smem_raw + 42496); + const int cluster_keys_addr = smem + 42496; + + // Mbarrier init (6 groups, 13 barriers) + // Mbarriers at smem_raw[0..104) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_ready: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // c_ready00: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_ready01: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_ready10: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_ready11: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // keys_ready: 8 barriers, init_count=8 + mbarrier_init_pred(smem + 40, 8, leader); + mbarrier_init_pred(smem + 48, 8, leader); + mbarrier_init_pred(smem + 56, 8, leader); + mbarrier_init_pred(smem + 64, 8, leader); + mbarrier_init_pred(smem + 72, 8, leader); + mbarrier_init_pred(smem + 80, 8, leader); + mbarrier_init_pred(smem + 88, 8, leader); + mbarrier_init_pred(smem + 96, 8, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); + + const int mbar_base = smem; + #define x_ready_addr (mbar_base + 0) + #define c_ready00_addr (mbar_base + 8) + #define c_ready01_addr (mbar_base + 16) + #define c_ready10_addr (mbar_base + 24) + #define c_ready11_addr (mbar_base + 32) + #define keys_ready_addr (mbar_base + 40) + + // === Task calls (dependency order) === + int total_tiles = B * num_n_tiles; + unsigned int _phase_x_ready_0 = 0; + unsigned int _phase_c_ready00_0 = 0; + unsigned int _phase_c_ready01_0 = 0; + unsigned int _phase_c_ready10_0 = 0; + unsigned int _phase_c_ready11_0 = 0; + unsigned int _phase_keys_ready_0 = 0; + #pragma unroll 1 + for (unsigned int tile = cluster_id; tile < total_tiles; tile += num_clusters) { + int batch = tile / (unsigned int)num_n_tiles; + int nt = tile % (unsigned int)num_n_tiles; + int point_base = batch * N + nt * 64; + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(x_ready_addr, 14336); + cp_async_bulk_gmem2smem(x_raw_addr, reinterpret_cast(reinterpret_cast(x) + ((unsigned long long)(point_base * D) * (unsigned long long)2)), 14336, x_ready_addr); + } + } + mbarrier_wait(x_ready_addr, _phase_x_ready_0); + _phase_x_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll 1 + for (unsigned int i = tid; i < 7168; i += 256) { + int row = i / 112; + int col = i % 112; + { + __nv_bfloat16 _bval_3404531696 = __float2bfloat16_rn(x_raw[i]); + uint16_t _bits_3404531696 = *(uint16_t*)&_bval_3404531696; + uint32_t _addr_3404531696 = static_cast((sx_addr + (unsigned int)(row * 224 + col * 2))); + asm volatile("st.shared.b16 [%0], %1;" :: "r"(_addr_3404531696), "h"(_bits_3404531696) : "memory"); + } + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + int warp_id_in_role = (warp - 0); + int group = warp_id_in_role / 4; + int local_warp = warp_id_in_role % 4; + int local_row = lane % 16; + int row_base = local_warp * 16; + float best = -3.4e+38f; + int owner_k_tiles = K / 64; + int group_k_tiles = owner_k_tiles / 2; + int group_tile_begin = cta_rank * owner_k_tiles + group * group_k_tiles; + int best_idx = group_tile_begin * 8; + int first_kbase = group_tile_begin * 8; + if (group == 0) { + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready00_addr, 1792); + cp_async_bulk_gmem2smem(c_direct00_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + first_kbase) * D) * (unsigned long long)2)), 1792, c_ready00_addr); + } + } + } else if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready10_addr, 1792); + cp_async_bulk_gmem2smem(c_direct10_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + first_kbase) * D) * (unsigned long long)2)), 1792, c_ready10_addr); + } + } + #pragma unroll 1 + for (int group_kt = 0; group_kt < group_k_tiles; group_kt++) { + int current_stage = group_kt % 2; + int kt = group_tile_begin + group_kt; + int kbase = kt * 8; + bool has_next = group_k_tiles > group_kt + 1; + if (has_next) { + int next_stage = (group_kt + 1) % 2; + int next_kbase = (kt + 1) * 8; + if (group == 0) { + if (next_stage == 0) { + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready00_addr, 1792); + cp_async_bulk_gmem2smem(c_direct00_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready00_addr); + } + } + } else if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready01_addr, 1792); + cp_async_bulk_gmem2smem(c_direct01_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready01_addr); + } + } + } else if (next_stage == 0) { + if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready10_addr, 1792); + cp_async_bulk_gmem2smem(c_direct10_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready10_addr); + } + } + } else { + if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready11_addr, 1792); + cp_async_bulk_gmem2smem(c_direct11_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready11_addr); + } + } + } + } + if (group == 0) { + if (current_stage == 0) { + mbarrier_wait(c_ready00_addr, _phase_c_ready00_0); + _phase_c_ready00_0 ^= 1; + } else { + mbarrier_wait(c_ready01_addr, _phase_c_ready01_0); + _phase_c_ready01_0 ^= 1; + } + } else if (current_stage == 0) { + mbarrier_wait(c_ready10_addr, _phase_c_ready10_0); + _phase_c_ready10_0 ^= 1; + } else { + mbarrier_wait(c_ready11_addr, _phase_c_ready11_0); + _phase_c_ready11_0 ^= 1; + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + unsigned int a[4]; + unsigned int b[2]; + float acc[4]; + unsigned int a_addr = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + lane / 16 * 16)); + unsigned int b_addr = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[0]), "=f"(acc[1]), "=f"(acc[2]), "=f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_0 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 32))); + unsigned int b_addr_1 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_0) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_1) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_2 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 64))); + unsigned int b_addr_3 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_2) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_3) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_4 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 96))); + unsigned int b_addr_5 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_4) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_5) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_6 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 128))); + unsigned int b_addr_7 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_6) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_7) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_8 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 160))); + unsigned int b_addr_9 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_8) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_9) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_10 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 192))); + unsigned int b_addr_11 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_10) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_11) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + #pragma unroll + for (int rp = 0; rp < 2; rp++) { + #pragma unroll + for (int cp = 0; cp < 2; cp++) { + int rr = lane / 4 + rp * 8; + int cc = lane % 4 * 2 + cp; + { + uint32_t _addr_3405188304 = static_cast((ss_addr + (unsigned int)((group * 64 + row_base + rr) * 32 + cc * 4))); + asm volatile("st.shared.f32 [%0], %1;" :: "r"(_addr_3405188304), "f"(acc[rp * 2 + cp]) : "memory"); + } + } + } + __syncwarp(); + if (lane < 16) { + #pragma unroll + for (int kk = 0; kk < 8; kk++) { + float score = ss[(group * 64 + row_base + local_row) * 8 + kk] - 0.5f * c_sq[batch * K + kbase + kk]; + if (score > best) { + best = score; + best_idx = kbase + kk; + } + } + } + if (group == 0) { + asm volatile("barrier.sync 1, %0;" :: "r"(128)); + } else { + asm volatile("barrier.sync 2, %0;" :: "r"(128)); + } + } + if (lane < 16) { + uint32_t _amf_u_0 = __float_as_uint(best); + uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; + unsigned int _amf_enc_0 = _amf_u_0 ^ _amf_mask_0; + unsigned long long shift64 = 32; + unsigned long long mask64 = 4294967295; + group_keys[group * 64 + row_base + local_row] = (unsigned long long)_amf_enc_0 << shift64 | mask64 - (unsigned long long)best_idx; + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + if (group == 0 && lane < 16) { + unsigned long long first_key = group_keys[row_base + local_row]; + unsigned long long second_key = group_keys[64 + row_base + local_row]; + local_keys[row_base + local_row] = ((second_key > first_key) ? second_key : first_key); + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + if (warp == 0) { + if (elect_sync()) { + uint32_t _mapa_0; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_0) : "r"(cluster_keys_addr + (unsigned int)(cta_rank * 512)), "r"(0)); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + uint32_t _mapa_1; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_1) : "r"(keys_ready_addr), "r"(0)); + asm volatile( + "mbarrier.arrive.expect_tx.release.cluster.shared::cluster.b64 _, [%0], %1;" + :: "r"(_mapa_1), "r"((uint32_t)(512)) : "memory"); + uint32_t _mapa_2; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_2) : "r"(keys_ready_addr), "r"(0)); + asm volatile( + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(_mapa_0), "r"(local_keys_addr), "r"((uint32_t)(512)), "r"(_mapa_2) + : "memory"); + } + } + if (cta_rank == 0) { + mbarrier_wait(keys_ready_addr, _phase_keys_ready_0); + _phase_keys_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int row_1 = tid / 2; + int lane_pair = tid % 2; + if (row_1 < 64) { + unsigned long long best_key = 0; + #pragma unroll + for (int peer = lane_pair; peer < 8; peer += 2) { + unsigned long long key = cluster_keys[peer * 64 + row_1]; + if (key > best_key) { + best_key = key; + } + } + unsigned long long _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best_key, 1); + unsigned long long peer_key = _shfl_xor_0; + if (peer_key > best_key) { + best_key = peer_key; + } + if (lane_pair == 0) { + unsigned long long mask64_1 = 4294967295; + int idx = (int)(mask64_1 - (best_key & mask64_1)); + *((int*)(out + (batch * N + nt * 64 + row_1))) = idx; + } + } + } + } + + // Cleanup + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0158.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0158.cu new file mode 100644 index 00000000..7ed658bb --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0158.cu @@ -0,0 +1,562 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_X_RAW_OFF 1024 +#define SMEM_X_RAW_STAGE_BYTES 14336 +#define SMEM_X_RAW_STRIDE 14336 +#define SMEM_C_DIRECT00_OFF 15360 +#define SMEM_C_DIRECT00_STAGE_BYTES 1792 +#define SMEM_C_DIRECT00_STRIDE 1792 +#define SMEM_C_DIRECT01_OFF 17152 +#define SMEM_C_DIRECT01_STAGE_BYTES 1792 +#define SMEM_C_DIRECT01_STRIDE 1792 +#define SMEM_C_DIRECT10_OFF 18944 +#define SMEM_C_DIRECT10_STAGE_BYTES 1792 +#define SMEM_C_DIRECT10_STRIDE 1792 +#define SMEM_C_DIRECT11_OFF 20736 +#define SMEM_C_DIRECT11_STAGE_BYTES 1792 +#define SMEM_C_DIRECT11_STRIDE 1792 +#define SMEM_SX_OFF 22528 +#define SMEM_SX_STAGE_BYTES 14336 +#define SMEM_SX_STRIDE 14336 +#define SMEM_SS_OFF 36864 +#define SMEM_SS_STAGE_BYTES 4096 +#define SMEM_SS_STRIDE 4096 +#define SMEM_GROUP_KEYS_OFF 40960 +#define SMEM_GROUP_KEYS_STAGE_BYTES 1024 +#define SMEM_GROUP_KEYS_STRIDE 1024 +#define SMEM_LOCAL_KEYS_OFF 41984 +#define SMEM_LOCAL_KEYS_STAGE_BYTES 512 +#define SMEM_LOCAL_KEYS_STRIDE 512 +#define SMEM_CLUSTER_KEYS_OFF 42496 +#define SMEM_CLUSTER_KEYS_STAGE_BYTES 4096 +#define SMEM_CLUSTER_KEYS_STRIDE 4096 +#define SMEM_TOTAL 46592 +#define THREADS 256 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ void cp_async_bulk_gmem2smem( + unsigned smem_addr, const void* gmem_ptr, unsigned bytes, int mbar_addr) { + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(smem_addr), "l"(gmem_ptr), "r"(bytes), "r"(mbar_addr) + : "memory"); +} + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 8; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 8; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* x_raw = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int x_raw_addr = smem + 1024; + __nv_bfloat16* c_direct00 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 15360); + const int c_direct00_addr = smem + 15360; + __nv_bfloat16* c_direct01 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17152); + const int c_direct01_addr = smem + 17152; + __nv_bfloat16* c_direct10 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 18944); + const int c_direct10_addr = smem + 18944; + __nv_bfloat16* c_direct11 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 20736); + const int c_direct11_addr = smem + 20736; + __nv_bfloat16* sx = reinterpret_cast<__nv_bfloat16*>(smem_raw + 22528); + const int sx_addr = smem + 22528; + float* ss = reinterpret_cast(smem_raw + 36864); + const int ss_addr = smem + 36864; + unsigned long long* group_keys = reinterpret_cast(smem_raw + 40960); + const int group_keys_addr = smem + 40960; + unsigned long long* local_keys = reinterpret_cast(smem_raw + 41984); + const int local_keys_addr = smem + 41984; + unsigned long long* cluster_keys = reinterpret_cast(smem_raw + 42496); + const int cluster_keys_addr = smem + 42496; + + // Mbarrier init (6 groups, 13 barriers) + // Mbarriers at smem_raw[0..104) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_ready: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // c_ready00: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_ready01: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_ready10: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_ready11: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // keys_ready: 8 barriers, init_count=8 + mbarrier_init_pred(smem + 40, 8, leader); + mbarrier_init_pred(smem + 48, 8, leader); + mbarrier_init_pred(smem + 56, 8, leader); + mbarrier_init_pred(smem + 64, 8, leader); + mbarrier_init_pred(smem + 72, 8, leader); + mbarrier_init_pred(smem + 80, 8, leader); + mbarrier_init_pred(smem + 88, 8, leader); + mbarrier_init_pred(smem + 96, 8, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); + + const int mbar_base = smem; + #define x_ready_addr (mbar_base + 0) + #define c_ready00_addr (mbar_base + 8) + #define c_ready01_addr (mbar_base + 16) + #define c_ready10_addr (mbar_base + 24) + #define c_ready11_addr (mbar_base + 32) + #define keys_ready_addr (mbar_base + 40) + + // === Task calls (dependency order) === + int total_tiles = B * num_n_tiles; + unsigned int _phase_x_ready_0 = 0; + unsigned int _phase_c_ready00_0 = 0; + unsigned int _phase_c_ready01_0 = 0; + unsigned int _phase_c_ready10_0 = 0; + unsigned int _phase_c_ready11_0 = 0; + unsigned int _phase_keys_ready_0 = 0; + #pragma unroll 1 + for (unsigned int tile = cluster_id; tile < total_tiles; tile += num_clusters) { + int batch = tile / (unsigned int)num_n_tiles; + int nt = tile % (unsigned int)num_n_tiles; + int point_base = batch * N + nt * 64; + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(x_ready_addr, 14336); + cp_async_bulk_gmem2smem(x_raw_addr, reinterpret_cast(reinterpret_cast(x) + ((unsigned long long)(point_base * D) * (unsigned long long)2)), 14336, x_ready_addr); + } + } + mbarrier_wait(x_ready_addr, _phase_x_ready_0); + _phase_x_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll 1 + for (unsigned int i = tid; i < 7168; i += 256) { + int row = i / 112; + int col = i % 112; + { + __nv_bfloat16 _bval_3404531696 = __float2bfloat16_rn(x_raw[i]); + uint16_t _bits_3404531696 = *(uint16_t*)&_bval_3404531696; + uint32_t _addr_3404531696 = static_cast((sx_addr + (unsigned int)(row * 224 + col * 2))); + asm volatile("st.shared.b16 [%0], %1;" :: "r"(_addr_3404531696), "h"(_bits_3404531696) : "memory"); + } + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + int warp_id_in_role = (warp - 0); + int group = warp_id_in_role / 4; + int local_warp = warp_id_in_role % 4; + int local_row = lane % 16; + int row_base = local_warp * 16; + float best = -3.4e+38f; + int owner_k_tiles = K / 64; + int group_k_tiles = owner_k_tiles / 2; + int group_tile_begin = cta_rank * owner_k_tiles + group * group_k_tiles; + int best_idx = group_tile_begin * 8; + int first_kbase = group_tile_begin * 8; + if (group == 0) { + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready00_addr, 1792); + cp_async_bulk_gmem2smem(c_direct00_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + first_kbase) * D) * (unsigned long long)2)), 1792, c_ready00_addr); + } + } + } else if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready10_addr, 1792); + cp_async_bulk_gmem2smem(c_direct10_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + first_kbase) * D) * (unsigned long long)2)), 1792, c_ready10_addr); + } + } + #pragma unroll 1 + for (int group_kt = 0; group_kt < group_k_tiles; group_kt++) { + int current_stage = group_kt % 2; + int kt = group_tile_begin + group_kt; + int kbase = kt * 8; + bool has_next = group_k_tiles > group_kt + 1; + if (has_next) { + int next_stage = (group_kt + 1) % 2; + int next_kbase = (kt + 1) * 8; + if (group == 0) { + if (next_stage == 0) { + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready00_addr, 1792); + cp_async_bulk_gmem2smem(c_direct00_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready00_addr); + } + } + } else if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready01_addr, 1792); + cp_async_bulk_gmem2smem(c_direct01_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready01_addr); + } + } + } else if (next_stage == 0) { + if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready10_addr, 1792); + cp_async_bulk_gmem2smem(c_direct10_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready10_addr); + } + } + } else { + if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready11_addr, 1792); + cp_async_bulk_gmem2smem(c_direct11_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready11_addr); + } + } + } + } + if (group == 0) { + if (current_stage == 0) { + mbarrier_wait(c_ready00_addr, _phase_c_ready00_0); + _phase_c_ready00_0 ^= 1; + } else { + mbarrier_wait(c_ready01_addr, _phase_c_ready01_0); + _phase_c_ready01_0 ^= 1; + } + } else if (current_stage == 0) { + mbarrier_wait(c_ready10_addr, _phase_c_ready10_0); + _phase_c_ready10_0 ^= 1; + } else { + mbarrier_wait(c_ready11_addr, _phase_c_ready11_0); + _phase_c_ready11_0 ^= 1; + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + unsigned int a[4]; + unsigned int b[2]; + float acc[4]; + unsigned int a_addr = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + lane / 16 * 16)); + unsigned int b_addr = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[0]), "=f"(acc[1]), "=f"(acc[2]), "=f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_0 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 32))); + unsigned int b_addr_1 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_0) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_1) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_2 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 64))); + unsigned int b_addr_3 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_2) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_3) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_4 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 96))); + unsigned int b_addr_5 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_4) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_5) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_6 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 128))); + unsigned int b_addr_7 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_6) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_7) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_8 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 160))); + unsigned int b_addr_9 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_8) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_9) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_10 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 192))); + unsigned int b_addr_11 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_10) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_11) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + #pragma unroll + for (int rp = 0; rp < 2; rp++) { + #pragma unroll + for (int cp = 0; cp < 2; cp++) { + int rr = lane / 4 + rp * 8; + int cc = lane % 4 * 2 + cp; + { + uint32_t _addr_3405188304 = static_cast((ss_addr + (unsigned int)((group * 64 + row_base + rr) * 32 + cc * 4))); + asm volatile("st.shared.f32 [%0], %1;" :: "r"(_addr_3405188304), "f"(acc[rp * 2 + cp]) : "memory"); + } + } + } + __syncwarp(); + if (lane < 16) { + #pragma unroll + for (int kk = 0; kk < 8; kk++) { + float score = ss[(group * 64 + row_base + local_row) * 8 + kk] - 0.5f * c_sq[batch * K + kbase + kk]; + if (score > best) { + best = score; + best_idx = kbase + kk; + } + } + } + if (group == 0) { + asm volatile("barrier.sync 1, %0;" :: "r"(128)); + } else { + asm volatile("barrier.sync 2, %0;" :: "r"(128)); + } + } + if (lane < 16) { + uint32_t _amf_u_0 = __float_as_uint(best); + uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; + unsigned int _amf_enc_0 = _amf_u_0 ^ _amf_mask_0; + unsigned long long shift64 = 32; + unsigned long long mask64 = 4294967295; + group_keys[group * 64 + row_base + local_row] = (unsigned long long)_amf_enc_0 << shift64 | mask64 - (unsigned long long)best_idx; + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + if (group == 0 && lane < 16) { + unsigned long long first_key = group_keys[row_base + local_row]; + unsigned long long second_key = group_keys[64 + row_base + local_row]; + local_keys[row_base + local_row] = ((second_key > first_key) ? second_key : first_key); + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + if (warp == 0) { + if (elect_sync()) { + uint32_t _mapa_0; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_0) : "r"(cluster_keys_addr + (unsigned int)(cta_rank * 512)), "r"(0)); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + uint32_t _mapa_1; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_1) : "r"(keys_ready_addr), "r"(0)); + asm volatile( + "mbarrier.arrive.expect_tx.release.cluster.shared::cluster.b64 _, [%0], %1;" + :: "r"(_mapa_1), "r"((uint32_t)(512)) : "memory"); + uint32_t _mapa_2; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_2) : "r"(keys_ready_addr), "r"(0)); + asm volatile( + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(_mapa_0), "r"(local_keys_addr), "r"((uint32_t)(512)), "r"(_mapa_2) + : "memory"); + } + } + if (cta_rank == 0) { + mbarrier_wait(keys_ready_addr, _phase_keys_ready_0); + _phase_keys_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int row_1 = tid / 2; + int lane_pair = tid % 2; + if (row_1 < 64) { + unsigned long long best_key = 0; + #pragma unroll + for (int peer = lane_pair; peer < 8; peer += 2) { + unsigned long long key = cluster_keys[peer * 64 + row_1]; + if (key > best_key) { + best_key = key; + } + } + unsigned long long _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best_key, 1); + unsigned long long peer_key = _shfl_xor_0; + if (peer_key > best_key) { + best_key = peer_key; + } + if (lane_pair == 0) { + unsigned long long mask64_1 = 4294967295; + int idx = (int)(mask64_1 - (best_key & mask64_1)); + *((int*)(out + (batch * N + nt * 64 + row_1))) = idx; + } + } + } + } + + // Cleanup + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0159.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0159.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0159.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0160.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0160.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0160.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0161.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0161.cu new file mode 100644 index 00000000..7ed658bb --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0161.cu @@ -0,0 +1,562 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_X_RAW_OFF 1024 +#define SMEM_X_RAW_STAGE_BYTES 14336 +#define SMEM_X_RAW_STRIDE 14336 +#define SMEM_C_DIRECT00_OFF 15360 +#define SMEM_C_DIRECT00_STAGE_BYTES 1792 +#define SMEM_C_DIRECT00_STRIDE 1792 +#define SMEM_C_DIRECT01_OFF 17152 +#define SMEM_C_DIRECT01_STAGE_BYTES 1792 +#define SMEM_C_DIRECT01_STRIDE 1792 +#define SMEM_C_DIRECT10_OFF 18944 +#define SMEM_C_DIRECT10_STAGE_BYTES 1792 +#define SMEM_C_DIRECT10_STRIDE 1792 +#define SMEM_C_DIRECT11_OFF 20736 +#define SMEM_C_DIRECT11_STAGE_BYTES 1792 +#define SMEM_C_DIRECT11_STRIDE 1792 +#define SMEM_SX_OFF 22528 +#define SMEM_SX_STAGE_BYTES 14336 +#define SMEM_SX_STRIDE 14336 +#define SMEM_SS_OFF 36864 +#define SMEM_SS_STAGE_BYTES 4096 +#define SMEM_SS_STRIDE 4096 +#define SMEM_GROUP_KEYS_OFF 40960 +#define SMEM_GROUP_KEYS_STAGE_BYTES 1024 +#define SMEM_GROUP_KEYS_STRIDE 1024 +#define SMEM_LOCAL_KEYS_OFF 41984 +#define SMEM_LOCAL_KEYS_STAGE_BYTES 512 +#define SMEM_LOCAL_KEYS_STRIDE 512 +#define SMEM_CLUSTER_KEYS_OFF 42496 +#define SMEM_CLUSTER_KEYS_STAGE_BYTES 4096 +#define SMEM_CLUSTER_KEYS_STRIDE 4096 +#define SMEM_TOTAL 46592 +#define THREADS 256 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ void cp_async_bulk_gmem2smem( + unsigned smem_addr, const void* gmem_ptr, unsigned bytes, int mbar_addr) { + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(smem_addr), "l"(gmem_ptr), "r"(bytes), "r"(mbar_addr) + : "memory"); +} + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 8; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 8; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* x_raw = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int x_raw_addr = smem + 1024; + __nv_bfloat16* c_direct00 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 15360); + const int c_direct00_addr = smem + 15360; + __nv_bfloat16* c_direct01 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17152); + const int c_direct01_addr = smem + 17152; + __nv_bfloat16* c_direct10 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 18944); + const int c_direct10_addr = smem + 18944; + __nv_bfloat16* c_direct11 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 20736); + const int c_direct11_addr = smem + 20736; + __nv_bfloat16* sx = reinterpret_cast<__nv_bfloat16*>(smem_raw + 22528); + const int sx_addr = smem + 22528; + float* ss = reinterpret_cast(smem_raw + 36864); + const int ss_addr = smem + 36864; + unsigned long long* group_keys = reinterpret_cast(smem_raw + 40960); + const int group_keys_addr = smem + 40960; + unsigned long long* local_keys = reinterpret_cast(smem_raw + 41984); + const int local_keys_addr = smem + 41984; + unsigned long long* cluster_keys = reinterpret_cast(smem_raw + 42496); + const int cluster_keys_addr = smem + 42496; + + // Mbarrier init (6 groups, 13 barriers) + // Mbarriers at smem_raw[0..104) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_ready: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // c_ready00: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_ready01: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_ready10: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_ready11: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // keys_ready: 8 barriers, init_count=8 + mbarrier_init_pred(smem + 40, 8, leader); + mbarrier_init_pred(smem + 48, 8, leader); + mbarrier_init_pred(smem + 56, 8, leader); + mbarrier_init_pred(smem + 64, 8, leader); + mbarrier_init_pred(smem + 72, 8, leader); + mbarrier_init_pred(smem + 80, 8, leader); + mbarrier_init_pred(smem + 88, 8, leader); + mbarrier_init_pred(smem + 96, 8, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); + + const int mbar_base = smem; + #define x_ready_addr (mbar_base + 0) + #define c_ready00_addr (mbar_base + 8) + #define c_ready01_addr (mbar_base + 16) + #define c_ready10_addr (mbar_base + 24) + #define c_ready11_addr (mbar_base + 32) + #define keys_ready_addr (mbar_base + 40) + + // === Task calls (dependency order) === + int total_tiles = B * num_n_tiles; + unsigned int _phase_x_ready_0 = 0; + unsigned int _phase_c_ready00_0 = 0; + unsigned int _phase_c_ready01_0 = 0; + unsigned int _phase_c_ready10_0 = 0; + unsigned int _phase_c_ready11_0 = 0; + unsigned int _phase_keys_ready_0 = 0; + #pragma unroll 1 + for (unsigned int tile = cluster_id; tile < total_tiles; tile += num_clusters) { + int batch = tile / (unsigned int)num_n_tiles; + int nt = tile % (unsigned int)num_n_tiles; + int point_base = batch * N + nt * 64; + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(x_ready_addr, 14336); + cp_async_bulk_gmem2smem(x_raw_addr, reinterpret_cast(reinterpret_cast(x) + ((unsigned long long)(point_base * D) * (unsigned long long)2)), 14336, x_ready_addr); + } + } + mbarrier_wait(x_ready_addr, _phase_x_ready_0); + _phase_x_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll 1 + for (unsigned int i = tid; i < 7168; i += 256) { + int row = i / 112; + int col = i % 112; + { + __nv_bfloat16 _bval_3404531696 = __float2bfloat16_rn(x_raw[i]); + uint16_t _bits_3404531696 = *(uint16_t*)&_bval_3404531696; + uint32_t _addr_3404531696 = static_cast((sx_addr + (unsigned int)(row * 224 + col * 2))); + asm volatile("st.shared.b16 [%0], %1;" :: "r"(_addr_3404531696), "h"(_bits_3404531696) : "memory"); + } + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + int warp_id_in_role = (warp - 0); + int group = warp_id_in_role / 4; + int local_warp = warp_id_in_role % 4; + int local_row = lane % 16; + int row_base = local_warp * 16; + float best = -3.4e+38f; + int owner_k_tiles = K / 64; + int group_k_tiles = owner_k_tiles / 2; + int group_tile_begin = cta_rank * owner_k_tiles + group * group_k_tiles; + int best_idx = group_tile_begin * 8; + int first_kbase = group_tile_begin * 8; + if (group == 0) { + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready00_addr, 1792); + cp_async_bulk_gmem2smem(c_direct00_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + first_kbase) * D) * (unsigned long long)2)), 1792, c_ready00_addr); + } + } + } else if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready10_addr, 1792); + cp_async_bulk_gmem2smem(c_direct10_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + first_kbase) * D) * (unsigned long long)2)), 1792, c_ready10_addr); + } + } + #pragma unroll 1 + for (int group_kt = 0; group_kt < group_k_tiles; group_kt++) { + int current_stage = group_kt % 2; + int kt = group_tile_begin + group_kt; + int kbase = kt * 8; + bool has_next = group_k_tiles > group_kt + 1; + if (has_next) { + int next_stage = (group_kt + 1) % 2; + int next_kbase = (kt + 1) * 8; + if (group == 0) { + if (next_stage == 0) { + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready00_addr, 1792); + cp_async_bulk_gmem2smem(c_direct00_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready00_addr); + } + } + } else if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready01_addr, 1792); + cp_async_bulk_gmem2smem(c_direct01_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready01_addr); + } + } + } else if (next_stage == 0) { + if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready10_addr, 1792); + cp_async_bulk_gmem2smem(c_direct10_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready10_addr); + } + } + } else { + if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready11_addr, 1792); + cp_async_bulk_gmem2smem(c_direct11_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready11_addr); + } + } + } + } + if (group == 0) { + if (current_stage == 0) { + mbarrier_wait(c_ready00_addr, _phase_c_ready00_0); + _phase_c_ready00_0 ^= 1; + } else { + mbarrier_wait(c_ready01_addr, _phase_c_ready01_0); + _phase_c_ready01_0 ^= 1; + } + } else if (current_stage == 0) { + mbarrier_wait(c_ready10_addr, _phase_c_ready10_0); + _phase_c_ready10_0 ^= 1; + } else { + mbarrier_wait(c_ready11_addr, _phase_c_ready11_0); + _phase_c_ready11_0 ^= 1; + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + unsigned int a[4]; + unsigned int b[2]; + float acc[4]; + unsigned int a_addr = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + lane / 16 * 16)); + unsigned int b_addr = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[0]), "=f"(acc[1]), "=f"(acc[2]), "=f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_0 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 32))); + unsigned int b_addr_1 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_0) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_1) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_2 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 64))); + unsigned int b_addr_3 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_2) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_3) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_4 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 96))); + unsigned int b_addr_5 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_4) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_5) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_6 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 128))); + unsigned int b_addr_7 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_6) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_7) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_8 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 160))); + unsigned int b_addr_9 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_8) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_9) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_10 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 192))); + unsigned int b_addr_11 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_10) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_11) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + #pragma unroll + for (int rp = 0; rp < 2; rp++) { + #pragma unroll + for (int cp = 0; cp < 2; cp++) { + int rr = lane / 4 + rp * 8; + int cc = lane % 4 * 2 + cp; + { + uint32_t _addr_3405188304 = static_cast((ss_addr + (unsigned int)((group * 64 + row_base + rr) * 32 + cc * 4))); + asm volatile("st.shared.f32 [%0], %1;" :: "r"(_addr_3405188304), "f"(acc[rp * 2 + cp]) : "memory"); + } + } + } + __syncwarp(); + if (lane < 16) { + #pragma unroll + for (int kk = 0; kk < 8; kk++) { + float score = ss[(group * 64 + row_base + local_row) * 8 + kk] - 0.5f * c_sq[batch * K + kbase + kk]; + if (score > best) { + best = score; + best_idx = kbase + kk; + } + } + } + if (group == 0) { + asm volatile("barrier.sync 1, %0;" :: "r"(128)); + } else { + asm volatile("barrier.sync 2, %0;" :: "r"(128)); + } + } + if (lane < 16) { + uint32_t _amf_u_0 = __float_as_uint(best); + uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; + unsigned int _amf_enc_0 = _amf_u_0 ^ _amf_mask_0; + unsigned long long shift64 = 32; + unsigned long long mask64 = 4294967295; + group_keys[group * 64 + row_base + local_row] = (unsigned long long)_amf_enc_0 << shift64 | mask64 - (unsigned long long)best_idx; + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + if (group == 0 && lane < 16) { + unsigned long long first_key = group_keys[row_base + local_row]; + unsigned long long second_key = group_keys[64 + row_base + local_row]; + local_keys[row_base + local_row] = ((second_key > first_key) ? second_key : first_key); + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + if (warp == 0) { + if (elect_sync()) { + uint32_t _mapa_0; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_0) : "r"(cluster_keys_addr + (unsigned int)(cta_rank * 512)), "r"(0)); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + uint32_t _mapa_1; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_1) : "r"(keys_ready_addr), "r"(0)); + asm volatile( + "mbarrier.arrive.expect_tx.release.cluster.shared::cluster.b64 _, [%0], %1;" + :: "r"(_mapa_1), "r"((uint32_t)(512)) : "memory"); + uint32_t _mapa_2; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_2) : "r"(keys_ready_addr), "r"(0)); + asm volatile( + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(_mapa_0), "r"(local_keys_addr), "r"((uint32_t)(512)), "r"(_mapa_2) + : "memory"); + } + } + if (cta_rank == 0) { + mbarrier_wait(keys_ready_addr, _phase_keys_ready_0); + _phase_keys_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int row_1 = tid / 2; + int lane_pair = tid % 2; + if (row_1 < 64) { + unsigned long long best_key = 0; + #pragma unroll + for (int peer = lane_pair; peer < 8; peer += 2) { + unsigned long long key = cluster_keys[peer * 64 + row_1]; + if (key > best_key) { + best_key = key; + } + } + unsigned long long _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best_key, 1); + unsigned long long peer_key = _shfl_xor_0; + if (peer_key > best_key) { + best_key = peer_key; + } + if (lane_pair == 0) { + unsigned long long mask64_1 = 4294967295; + int idx = (int)(mask64_1 - (best_key & mask64_1)); + *((int*)(out + (batch * N + nt * 64 + row_1))) = idx; + } + } + } + } + + // Cleanup + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0162.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0162.cu new file mode 100644 index 00000000..7ed658bb --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0162.cu @@ -0,0 +1,562 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_X_RAW_OFF 1024 +#define SMEM_X_RAW_STAGE_BYTES 14336 +#define SMEM_X_RAW_STRIDE 14336 +#define SMEM_C_DIRECT00_OFF 15360 +#define SMEM_C_DIRECT00_STAGE_BYTES 1792 +#define SMEM_C_DIRECT00_STRIDE 1792 +#define SMEM_C_DIRECT01_OFF 17152 +#define SMEM_C_DIRECT01_STAGE_BYTES 1792 +#define SMEM_C_DIRECT01_STRIDE 1792 +#define SMEM_C_DIRECT10_OFF 18944 +#define SMEM_C_DIRECT10_STAGE_BYTES 1792 +#define SMEM_C_DIRECT10_STRIDE 1792 +#define SMEM_C_DIRECT11_OFF 20736 +#define SMEM_C_DIRECT11_STAGE_BYTES 1792 +#define SMEM_C_DIRECT11_STRIDE 1792 +#define SMEM_SX_OFF 22528 +#define SMEM_SX_STAGE_BYTES 14336 +#define SMEM_SX_STRIDE 14336 +#define SMEM_SS_OFF 36864 +#define SMEM_SS_STAGE_BYTES 4096 +#define SMEM_SS_STRIDE 4096 +#define SMEM_GROUP_KEYS_OFF 40960 +#define SMEM_GROUP_KEYS_STAGE_BYTES 1024 +#define SMEM_GROUP_KEYS_STRIDE 1024 +#define SMEM_LOCAL_KEYS_OFF 41984 +#define SMEM_LOCAL_KEYS_STAGE_BYTES 512 +#define SMEM_LOCAL_KEYS_STRIDE 512 +#define SMEM_CLUSTER_KEYS_OFF 42496 +#define SMEM_CLUSTER_KEYS_STAGE_BYTES 4096 +#define SMEM_CLUSTER_KEYS_STRIDE 4096 +#define SMEM_TOTAL 46592 +#define THREADS 256 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ void cp_async_bulk_gmem2smem( + unsigned smem_addr, const void* gmem_ptr, unsigned bytes, int mbar_addr) { + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(smem_addr), "l"(gmem_ptr), "r"(bytes), "r"(mbar_addr) + : "memory"); +} + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 8; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 8; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* x_raw = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int x_raw_addr = smem + 1024; + __nv_bfloat16* c_direct00 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 15360); + const int c_direct00_addr = smem + 15360; + __nv_bfloat16* c_direct01 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17152); + const int c_direct01_addr = smem + 17152; + __nv_bfloat16* c_direct10 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 18944); + const int c_direct10_addr = smem + 18944; + __nv_bfloat16* c_direct11 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 20736); + const int c_direct11_addr = smem + 20736; + __nv_bfloat16* sx = reinterpret_cast<__nv_bfloat16*>(smem_raw + 22528); + const int sx_addr = smem + 22528; + float* ss = reinterpret_cast(smem_raw + 36864); + const int ss_addr = smem + 36864; + unsigned long long* group_keys = reinterpret_cast(smem_raw + 40960); + const int group_keys_addr = smem + 40960; + unsigned long long* local_keys = reinterpret_cast(smem_raw + 41984); + const int local_keys_addr = smem + 41984; + unsigned long long* cluster_keys = reinterpret_cast(smem_raw + 42496); + const int cluster_keys_addr = smem + 42496; + + // Mbarrier init (6 groups, 13 barriers) + // Mbarriers at smem_raw[0..104) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_ready: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // c_ready00: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_ready01: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_ready10: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_ready11: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // keys_ready: 8 barriers, init_count=8 + mbarrier_init_pred(smem + 40, 8, leader); + mbarrier_init_pred(smem + 48, 8, leader); + mbarrier_init_pred(smem + 56, 8, leader); + mbarrier_init_pred(smem + 64, 8, leader); + mbarrier_init_pred(smem + 72, 8, leader); + mbarrier_init_pred(smem + 80, 8, leader); + mbarrier_init_pred(smem + 88, 8, leader); + mbarrier_init_pred(smem + 96, 8, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); + + const int mbar_base = smem; + #define x_ready_addr (mbar_base + 0) + #define c_ready00_addr (mbar_base + 8) + #define c_ready01_addr (mbar_base + 16) + #define c_ready10_addr (mbar_base + 24) + #define c_ready11_addr (mbar_base + 32) + #define keys_ready_addr (mbar_base + 40) + + // === Task calls (dependency order) === + int total_tiles = B * num_n_tiles; + unsigned int _phase_x_ready_0 = 0; + unsigned int _phase_c_ready00_0 = 0; + unsigned int _phase_c_ready01_0 = 0; + unsigned int _phase_c_ready10_0 = 0; + unsigned int _phase_c_ready11_0 = 0; + unsigned int _phase_keys_ready_0 = 0; + #pragma unroll 1 + for (unsigned int tile = cluster_id; tile < total_tiles; tile += num_clusters) { + int batch = tile / (unsigned int)num_n_tiles; + int nt = tile % (unsigned int)num_n_tiles; + int point_base = batch * N + nt * 64; + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(x_ready_addr, 14336); + cp_async_bulk_gmem2smem(x_raw_addr, reinterpret_cast(reinterpret_cast(x) + ((unsigned long long)(point_base * D) * (unsigned long long)2)), 14336, x_ready_addr); + } + } + mbarrier_wait(x_ready_addr, _phase_x_ready_0); + _phase_x_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll 1 + for (unsigned int i = tid; i < 7168; i += 256) { + int row = i / 112; + int col = i % 112; + { + __nv_bfloat16 _bval_3404531696 = __float2bfloat16_rn(x_raw[i]); + uint16_t _bits_3404531696 = *(uint16_t*)&_bval_3404531696; + uint32_t _addr_3404531696 = static_cast((sx_addr + (unsigned int)(row * 224 + col * 2))); + asm volatile("st.shared.b16 [%0], %1;" :: "r"(_addr_3404531696), "h"(_bits_3404531696) : "memory"); + } + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + int warp_id_in_role = (warp - 0); + int group = warp_id_in_role / 4; + int local_warp = warp_id_in_role % 4; + int local_row = lane % 16; + int row_base = local_warp * 16; + float best = -3.4e+38f; + int owner_k_tiles = K / 64; + int group_k_tiles = owner_k_tiles / 2; + int group_tile_begin = cta_rank * owner_k_tiles + group * group_k_tiles; + int best_idx = group_tile_begin * 8; + int first_kbase = group_tile_begin * 8; + if (group == 0) { + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready00_addr, 1792); + cp_async_bulk_gmem2smem(c_direct00_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + first_kbase) * D) * (unsigned long long)2)), 1792, c_ready00_addr); + } + } + } else if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready10_addr, 1792); + cp_async_bulk_gmem2smem(c_direct10_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + first_kbase) * D) * (unsigned long long)2)), 1792, c_ready10_addr); + } + } + #pragma unroll 1 + for (int group_kt = 0; group_kt < group_k_tiles; group_kt++) { + int current_stage = group_kt % 2; + int kt = group_tile_begin + group_kt; + int kbase = kt * 8; + bool has_next = group_k_tiles > group_kt + 1; + if (has_next) { + int next_stage = (group_kt + 1) % 2; + int next_kbase = (kt + 1) * 8; + if (group == 0) { + if (next_stage == 0) { + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready00_addr, 1792); + cp_async_bulk_gmem2smem(c_direct00_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready00_addr); + } + } + } else if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready01_addr, 1792); + cp_async_bulk_gmem2smem(c_direct01_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready01_addr); + } + } + } else if (next_stage == 0) { + if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready10_addr, 1792); + cp_async_bulk_gmem2smem(c_direct10_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready10_addr); + } + } + } else { + if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready11_addr, 1792); + cp_async_bulk_gmem2smem(c_direct11_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready11_addr); + } + } + } + } + if (group == 0) { + if (current_stage == 0) { + mbarrier_wait(c_ready00_addr, _phase_c_ready00_0); + _phase_c_ready00_0 ^= 1; + } else { + mbarrier_wait(c_ready01_addr, _phase_c_ready01_0); + _phase_c_ready01_0 ^= 1; + } + } else if (current_stage == 0) { + mbarrier_wait(c_ready10_addr, _phase_c_ready10_0); + _phase_c_ready10_0 ^= 1; + } else { + mbarrier_wait(c_ready11_addr, _phase_c_ready11_0); + _phase_c_ready11_0 ^= 1; + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + unsigned int a[4]; + unsigned int b[2]; + float acc[4]; + unsigned int a_addr = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + lane / 16 * 16)); + unsigned int b_addr = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[0]), "=f"(acc[1]), "=f"(acc[2]), "=f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_0 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 32))); + unsigned int b_addr_1 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_0) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_1) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_2 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 64))); + unsigned int b_addr_3 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_2) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_3) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_4 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 96))); + unsigned int b_addr_5 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_4) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_5) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_6 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 128))); + unsigned int b_addr_7 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_6) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_7) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_8 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 160))); + unsigned int b_addr_9 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_8) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_9) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_10 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 192))); + unsigned int b_addr_11 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_10) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_11) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + #pragma unroll + for (int rp = 0; rp < 2; rp++) { + #pragma unroll + for (int cp = 0; cp < 2; cp++) { + int rr = lane / 4 + rp * 8; + int cc = lane % 4 * 2 + cp; + { + uint32_t _addr_3405188304 = static_cast((ss_addr + (unsigned int)((group * 64 + row_base + rr) * 32 + cc * 4))); + asm volatile("st.shared.f32 [%0], %1;" :: "r"(_addr_3405188304), "f"(acc[rp * 2 + cp]) : "memory"); + } + } + } + __syncwarp(); + if (lane < 16) { + #pragma unroll + for (int kk = 0; kk < 8; kk++) { + float score = ss[(group * 64 + row_base + local_row) * 8 + kk] - 0.5f * c_sq[batch * K + kbase + kk]; + if (score > best) { + best = score; + best_idx = kbase + kk; + } + } + } + if (group == 0) { + asm volatile("barrier.sync 1, %0;" :: "r"(128)); + } else { + asm volatile("barrier.sync 2, %0;" :: "r"(128)); + } + } + if (lane < 16) { + uint32_t _amf_u_0 = __float_as_uint(best); + uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; + unsigned int _amf_enc_0 = _amf_u_0 ^ _amf_mask_0; + unsigned long long shift64 = 32; + unsigned long long mask64 = 4294967295; + group_keys[group * 64 + row_base + local_row] = (unsigned long long)_amf_enc_0 << shift64 | mask64 - (unsigned long long)best_idx; + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + if (group == 0 && lane < 16) { + unsigned long long first_key = group_keys[row_base + local_row]; + unsigned long long second_key = group_keys[64 + row_base + local_row]; + local_keys[row_base + local_row] = ((second_key > first_key) ? second_key : first_key); + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + if (warp == 0) { + if (elect_sync()) { + uint32_t _mapa_0; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_0) : "r"(cluster_keys_addr + (unsigned int)(cta_rank * 512)), "r"(0)); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + uint32_t _mapa_1; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_1) : "r"(keys_ready_addr), "r"(0)); + asm volatile( + "mbarrier.arrive.expect_tx.release.cluster.shared::cluster.b64 _, [%0], %1;" + :: "r"(_mapa_1), "r"((uint32_t)(512)) : "memory"); + uint32_t _mapa_2; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_2) : "r"(keys_ready_addr), "r"(0)); + asm volatile( + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(_mapa_0), "r"(local_keys_addr), "r"((uint32_t)(512)), "r"(_mapa_2) + : "memory"); + } + } + if (cta_rank == 0) { + mbarrier_wait(keys_ready_addr, _phase_keys_ready_0); + _phase_keys_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int row_1 = tid / 2; + int lane_pair = tid % 2; + if (row_1 < 64) { + unsigned long long best_key = 0; + #pragma unroll + for (int peer = lane_pair; peer < 8; peer += 2) { + unsigned long long key = cluster_keys[peer * 64 + row_1]; + if (key > best_key) { + best_key = key; + } + } + unsigned long long _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best_key, 1); + unsigned long long peer_key = _shfl_xor_0; + if (peer_key > best_key) { + best_key = peer_key; + } + if (lane_pair == 0) { + unsigned long long mask64_1 = 4294967295; + int idx = (int)(mask64_1 - (best_key & mask64_1)); + *((int*)(out + (batch * N + nt * 64 + row_1))) = idx; + } + } + } + } + + // Cleanup + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0163.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0163.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0163.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0164.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0164.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0164.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0165.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0165.cu new file mode 100644 index 00000000..7ed658bb --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0165.cu @@ -0,0 +1,562 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_X_RAW_OFF 1024 +#define SMEM_X_RAW_STAGE_BYTES 14336 +#define SMEM_X_RAW_STRIDE 14336 +#define SMEM_C_DIRECT00_OFF 15360 +#define SMEM_C_DIRECT00_STAGE_BYTES 1792 +#define SMEM_C_DIRECT00_STRIDE 1792 +#define SMEM_C_DIRECT01_OFF 17152 +#define SMEM_C_DIRECT01_STAGE_BYTES 1792 +#define SMEM_C_DIRECT01_STRIDE 1792 +#define SMEM_C_DIRECT10_OFF 18944 +#define SMEM_C_DIRECT10_STAGE_BYTES 1792 +#define SMEM_C_DIRECT10_STRIDE 1792 +#define SMEM_C_DIRECT11_OFF 20736 +#define SMEM_C_DIRECT11_STAGE_BYTES 1792 +#define SMEM_C_DIRECT11_STRIDE 1792 +#define SMEM_SX_OFF 22528 +#define SMEM_SX_STAGE_BYTES 14336 +#define SMEM_SX_STRIDE 14336 +#define SMEM_SS_OFF 36864 +#define SMEM_SS_STAGE_BYTES 4096 +#define SMEM_SS_STRIDE 4096 +#define SMEM_GROUP_KEYS_OFF 40960 +#define SMEM_GROUP_KEYS_STAGE_BYTES 1024 +#define SMEM_GROUP_KEYS_STRIDE 1024 +#define SMEM_LOCAL_KEYS_OFF 41984 +#define SMEM_LOCAL_KEYS_STAGE_BYTES 512 +#define SMEM_LOCAL_KEYS_STRIDE 512 +#define SMEM_CLUSTER_KEYS_OFF 42496 +#define SMEM_CLUSTER_KEYS_STAGE_BYTES 4096 +#define SMEM_CLUSTER_KEYS_STRIDE 4096 +#define SMEM_TOTAL 46592 +#define THREADS 256 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ void cp_async_bulk_gmem2smem( + unsigned smem_addr, const void* gmem_ptr, unsigned bytes, int mbar_addr) { + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(smem_addr), "l"(gmem_ptr), "r"(bytes), "r"(mbar_addr) + : "memory"); +} + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 8; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 8; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* x_raw = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int x_raw_addr = smem + 1024; + __nv_bfloat16* c_direct00 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 15360); + const int c_direct00_addr = smem + 15360; + __nv_bfloat16* c_direct01 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17152); + const int c_direct01_addr = smem + 17152; + __nv_bfloat16* c_direct10 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 18944); + const int c_direct10_addr = smem + 18944; + __nv_bfloat16* c_direct11 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 20736); + const int c_direct11_addr = smem + 20736; + __nv_bfloat16* sx = reinterpret_cast<__nv_bfloat16*>(smem_raw + 22528); + const int sx_addr = smem + 22528; + float* ss = reinterpret_cast(smem_raw + 36864); + const int ss_addr = smem + 36864; + unsigned long long* group_keys = reinterpret_cast(smem_raw + 40960); + const int group_keys_addr = smem + 40960; + unsigned long long* local_keys = reinterpret_cast(smem_raw + 41984); + const int local_keys_addr = smem + 41984; + unsigned long long* cluster_keys = reinterpret_cast(smem_raw + 42496); + const int cluster_keys_addr = smem + 42496; + + // Mbarrier init (6 groups, 13 barriers) + // Mbarriers at smem_raw[0..104) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_ready: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // c_ready00: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_ready01: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_ready10: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_ready11: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // keys_ready: 8 barriers, init_count=8 + mbarrier_init_pred(smem + 40, 8, leader); + mbarrier_init_pred(smem + 48, 8, leader); + mbarrier_init_pred(smem + 56, 8, leader); + mbarrier_init_pred(smem + 64, 8, leader); + mbarrier_init_pred(smem + 72, 8, leader); + mbarrier_init_pred(smem + 80, 8, leader); + mbarrier_init_pred(smem + 88, 8, leader); + mbarrier_init_pred(smem + 96, 8, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); + + const int mbar_base = smem; + #define x_ready_addr (mbar_base + 0) + #define c_ready00_addr (mbar_base + 8) + #define c_ready01_addr (mbar_base + 16) + #define c_ready10_addr (mbar_base + 24) + #define c_ready11_addr (mbar_base + 32) + #define keys_ready_addr (mbar_base + 40) + + // === Task calls (dependency order) === + int total_tiles = B * num_n_tiles; + unsigned int _phase_x_ready_0 = 0; + unsigned int _phase_c_ready00_0 = 0; + unsigned int _phase_c_ready01_0 = 0; + unsigned int _phase_c_ready10_0 = 0; + unsigned int _phase_c_ready11_0 = 0; + unsigned int _phase_keys_ready_0 = 0; + #pragma unroll 1 + for (unsigned int tile = cluster_id; tile < total_tiles; tile += num_clusters) { + int batch = tile / (unsigned int)num_n_tiles; + int nt = tile % (unsigned int)num_n_tiles; + int point_base = batch * N + nt * 64; + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(x_ready_addr, 14336); + cp_async_bulk_gmem2smem(x_raw_addr, reinterpret_cast(reinterpret_cast(x) + ((unsigned long long)(point_base * D) * (unsigned long long)2)), 14336, x_ready_addr); + } + } + mbarrier_wait(x_ready_addr, _phase_x_ready_0); + _phase_x_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll 1 + for (unsigned int i = tid; i < 7168; i += 256) { + int row = i / 112; + int col = i % 112; + { + __nv_bfloat16 _bval_3404531696 = __float2bfloat16_rn(x_raw[i]); + uint16_t _bits_3404531696 = *(uint16_t*)&_bval_3404531696; + uint32_t _addr_3404531696 = static_cast((sx_addr + (unsigned int)(row * 224 + col * 2))); + asm volatile("st.shared.b16 [%0], %1;" :: "r"(_addr_3404531696), "h"(_bits_3404531696) : "memory"); + } + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + int warp_id_in_role = (warp - 0); + int group = warp_id_in_role / 4; + int local_warp = warp_id_in_role % 4; + int local_row = lane % 16; + int row_base = local_warp * 16; + float best = -3.4e+38f; + int owner_k_tiles = K / 64; + int group_k_tiles = owner_k_tiles / 2; + int group_tile_begin = cta_rank * owner_k_tiles + group * group_k_tiles; + int best_idx = group_tile_begin * 8; + int first_kbase = group_tile_begin * 8; + if (group == 0) { + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready00_addr, 1792); + cp_async_bulk_gmem2smem(c_direct00_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + first_kbase) * D) * (unsigned long long)2)), 1792, c_ready00_addr); + } + } + } else if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready10_addr, 1792); + cp_async_bulk_gmem2smem(c_direct10_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + first_kbase) * D) * (unsigned long long)2)), 1792, c_ready10_addr); + } + } + #pragma unroll 1 + for (int group_kt = 0; group_kt < group_k_tiles; group_kt++) { + int current_stage = group_kt % 2; + int kt = group_tile_begin + group_kt; + int kbase = kt * 8; + bool has_next = group_k_tiles > group_kt + 1; + if (has_next) { + int next_stage = (group_kt + 1) % 2; + int next_kbase = (kt + 1) * 8; + if (group == 0) { + if (next_stage == 0) { + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready00_addr, 1792); + cp_async_bulk_gmem2smem(c_direct00_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready00_addr); + } + } + } else if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready01_addr, 1792); + cp_async_bulk_gmem2smem(c_direct01_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready01_addr); + } + } + } else if (next_stage == 0) { + if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready10_addr, 1792); + cp_async_bulk_gmem2smem(c_direct10_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready10_addr); + } + } + } else { + if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready11_addr, 1792); + cp_async_bulk_gmem2smem(c_direct11_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready11_addr); + } + } + } + } + if (group == 0) { + if (current_stage == 0) { + mbarrier_wait(c_ready00_addr, _phase_c_ready00_0); + _phase_c_ready00_0 ^= 1; + } else { + mbarrier_wait(c_ready01_addr, _phase_c_ready01_0); + _phase_c_ready01_0 ^= 1; + } + } else if (current_stage == 0) { + mbarrier_wait(c_ready10_addr, _phase_c_ready10_0); + _phase_c_ready10_0 ^= 1; + } else { + mbarrier_wait(c_ready11_addr, _phase_c_ready11_0); + _phase_c_ready11_0 ^= 1; + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + unsigned int a[4]; + unsigned int b[2]; + float acc[4]; + unsigned int a_addr = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + lane / 16 * 16)); + unsigned int b_addr = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[0]), "=f"(acc[1]), "=f"(acc[2]), "=f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_0 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 32))); + unsigned int b_addr_1 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_0) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_1) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_2 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 64))); + unsigned int b_addr_3 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_2) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_3) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_4 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 96))); + unsigned int b_addr_5 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_4) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_5) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_6 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 128))); + unsigned int b_addr_7 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_6) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_7) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_8 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 160))); + unsigned int b_addr_9 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_8) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_9) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_10 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 192))); + unsigned int b_addr_11 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_10) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_11) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + #pragma unroll + for (int rp = 0; rp < 2; rp++) { + #pragma unroll + for (int cp = 0; cp < 2; cp++) { + int rr = lane / 4 + rp * 8; + int cc = lane % 4 * 2 + cp; + { + uint32_t _addr_3405188304 = static_cast((ss_addr + (unsigned int)((group * 64 + row_base + rr) * 32 + cc * 4))); + asm volatile("st.shared.f32 [%0], %1;" :: "r"(_addr_3405188304), "f"(acc[rp * 2 + cp]) : "memory"); + } + } + } + __syncwarp(); + if (lane < 16) { + #pragma unroll + for (int kk = 0; kk < 8; kk++) { + float score = ss[(group * 64 + row_base + local_row) * 8 + kk] - 0.5f * c_sq[batch * K + kbase + kk]; + if (score > best) { + best = score; + best_idx = kbase + kk; + } + } + } + if (group == 0) { + asm volatile("barrier.sync 1, %0;" :: "r"(128)); + } else { + asm volatile("barrier.sync 2, %0;" :: "r"(128)); + } + } + if (lane < 16) { + uint32_t _amf_u_0 = __float_as_uint(best); + uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; + unsigned int _amf_enc_0 = _amf_u_0 ^ _amf_mask_0; + unsigned long long shift64 = 32; + unsigned long long mask64 = 4294967295; + group_keys[group * 64 + row_base + local_row] = (unsigned long long)_amf_enc_0 << shift64 | mask64 - (unsigned long long)best_idx; + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + if (group == 0 && lane < 16) { + unsigned long long first_key = group_keys[row_base + local_row]; + unsigned long long second_key = group_keys[64 + row_base + local_row]; + local_keys[row_base + local_row] = ((second_key > first_key) ? second_key : first_key); + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + if (warp == 0) { + if (elect_sync()) { + uint32_t _mapa_0; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_0) : "r"(cluster_keys_addr + (unsigned int)(cta_rank * 512)), "r"(0)); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + uint32_t _mapa_1; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_1) : "r"(keys_ready_addr), "r"(0)); + asm volatile( + "mbarrier.arrive.expect_tx.release.cluster.shared::cluster.b64 _, [%0], %1;" + :: "r"(_mapa_1), "r"((uint32_t)(512)) : "memory"); + uint32_t _mapa_2; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_2) : "r"(keys_ready_addr), "r"(0)); + asm volatile( + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(_mapa_0), "r"(local_keys_addr), "r"((uint32_t)(512)), "r"(_mapa_2) + : "memory"); + } + } + if (cta_rank == 0) { + mbarrier_wait(keys_ready_addr, _phase_keys_ready_0); + _phase_keys_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int row_1 = tid / 2; + int lane_pair = tid % 2; + if (row_1 < 64) { + unsigned long long best_key = 0; + #pragma unroll + for (int peer = lane_pair; peer < 8; peer += 2) { + unsigned long long key = cluster_keys[peer * 64 + row_1]; + if (key > best_key) { + best_key = key; + } + } + unsigned long long _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best_key, 1); + unsigned long long peer_key = _shfl_xor_0; + if (peer_key > best_key) { + best_key = peer_key; + } + if (lane_pair == 0) { + unsigned long long mask64_1 = 4294967295; + int idx = (int)(mask64_1 - (best_key & mask64_1)); + *((int*)(out + (batch * N + nt * 64 + row_1))) = idx; + } + } + } + } + + // Cleanup + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0166.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0166.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0166.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0167.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0167.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0167.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0168.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0168.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0168.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0169.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0169.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0169.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0170.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0170.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0170.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0171.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0171.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0171.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0172.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0172.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0172.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0173.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0173.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0173.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0174.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0174.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0174.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0175.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0175.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0175.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0176.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0176.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0176.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0177.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0177.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0177.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0178.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0178.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0178.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0179.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0179.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0179.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0180.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0180.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0180.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0181.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0181.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0181.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0182.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0182.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0182.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0183.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0183.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0183.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0184.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0184.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0184.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0185.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0185.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0185.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0186.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0186.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0186.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0187.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0187.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0187.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0188.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0188.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0188.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0189.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0189.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0189.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0190.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0190.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0190.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0191.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0191.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0191.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0192.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0192.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0192.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0193.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0193.cu new file mode 100644 index 00000000..482ca687 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0193.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 24 * 8; + int row = vec_idx / 24; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 24 * 8; + int row_1 = vec_idx_1 / 24; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0194.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0194.cu new file mode 100644 index 00000000..7070fa90 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0194.cu @@ -0,0 +1,600 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 49152 +#define SMEM_SMEM_X_STRIDE 49152 +#define SMEM_SMEM_C_OFF 50176 +#define SMEM_SMEM_C_STAGE_BYTES 98304 +#define SMEM_SMEM_C_STRIDE 98304 +#define SMEM_SMEM_CSQ_OFF 148480 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 149504 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int smem_c_addr = smem + 50176; + float* smem_csq = reinterpret_cast(smem_raw + 148480); + const int smem_csq_addr = smem + 148480; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 49152); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 98304); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0195.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0195.cu new file mode 100644 index 00000000..482ca687 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0195.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 24 * 8; + int row = vec_idx / 24; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 24 * 8; + int row_1 = vec_idx_1 / 24; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0196.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0196.cu new file mode 100644 index 00000000..17caae28 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0196.cu @@ -0,0 +1,1211 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 49152 +#define SMEM_SMEM_X0_STRIDE 49152 +#define SMEM_SMEM_X1_OFF 50176 +#define SMEM_SMEM_X1_STAGE_BYTES 49152 +#define SMEM_SMEM_X1_STRIDE 49152 +#define SMEM_SMEM_C_OFF 99328 +#define SMEM_SMEM_C_STAGE_BYTES 98304 +#define SMEM_SMEM_C_STRIDE 98304 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int smem_x1_addr = smem + 50176; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); + const int smem_c_addr = smem + 99328; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 49152); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 49152); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 98304); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x0_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_4 = smem_x1_addr + 16384; + int _mma_a_lo_4 = make_warp_uniform((_mma_a_addr_4 >> 4) & 0x3FFF); + int _mma_b_addr_4 = smem_c_addr + 32768; + int _mma_b_lo_4 = make_warp_uniform((_mma_b_addr_4 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_5 = smem_x1_addr + 32768; + int _mma_a_lo_5 = make_warp_uniform((_mma_a_addr_5 >> 4) & 0x3FFF); + int _mma_b_addr_5 = smem_c_addr + 65536; + int _mma_b_lo_5 = make_warp_uniform((_mma_b_addr_5 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0197.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0197.cu new file mode 100644 index 00000000..482ca687 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0197.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 24 * 8; + int row = vec_idx / 24; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 24 * 8; + int row_1 = vec_idx_1 / 24; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0198.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0198.cu new file mode 100644 index 00000000..17caae28 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0198.cu @@ -0,0 +1,1211 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 49152 +#define SMEM_SMEM_X0_STRIDE 49152 +#define SMEM_SMEM_X1_OFF 50176 +#define SMEM_SMEM_X1_STAGE_BYTES 49152 +#define SMEM_SMEM_X1_STRIDE 49152 +#define SMEM_SMEM_C_OFF 99328 +#define SMEM_SMEM_C_STAGE_BYTES 98304 +#define SMEM_SMEM_C_STRIDE 98304 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int smem_x1_addr = smem + 50176; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); + const int smem_c_addr = smem + 99328; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 49152); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 49152); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 98304); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x0_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_4 = smem_x1_addr + 16384; + int _mma_a_lo_4 = make_warp_uniform((_mma_a_addr_4 >> 4) & 0x3FFF); + int _mma_b_addr_4 = smem_c_addr + 32768; + int _mma_b_lo_4 = make_warp_uniform((_mma_b_addr_4 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_5 = smem_x1_addr + 32768; + int _mma_a_lo_5 = make_warp_uniform((_mma_a_addr_5 >> 4) & 0x3FFF); + int _mma_b_addr_5 = smem_c_addr + 65536; + int _mma_b_lo_5 = make_warp_uniform((_mma_b_addr_5 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0199.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0199.cu new file mode 100644 index 00000000..482ca687 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0199.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 24 * 8; + int row = vec_idx / 24; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 24 * 8; + int row_1 = vec_idx_1 / 24; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0200.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0200.cu new file mode 100644 index 00000000..17caae28 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0200.cu @@ -0,0 +1,1211 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 49152 +#define SMEM_SMEM_X0_STRIDE 49152 +#define SMEM_SMEM_X1_OFF 50176 +#define SMEM_SMEM_X1_STAGE_BYTES 49152 +#define SMEM_SMEM_X1_STRIDE 49152 +#define SMEM_SMEM_C_OFF 99328 +#define SMEM_SMEM_C_STAGE_BYTES 98304 +#define SMEM_SMEM_C_STRIDE 98304 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int smem_x1_addr = smem + 50176; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); + const int smem_c_addr = smem + 99328; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 49152); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 49152); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 98304); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x0_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_4 = smem_x1_addr + 16384; + int _mma_a_lo_4 = make_warp_uniform((_mma_a_addr_4 >> 4) & 0x3FFF); + int _mma_b_addr_4 = smem_c_addr + 32768; + int _mma_b_lo_4 = make_warp_uniform((_mma_b_addr_4 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_5 = smem_x1_addr + 32768; + int _mma_a_lo_5 = make_warp_uniform((_mma_a_addr_5 >> 4) & 0x3FFF); + int _mma_b_addr_5 = smem_c_addr + 65536; + int _mma_b_lo_5 = make_warp_uniform((_mma_b_addr_5 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0201.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0201.cu new file mode 100644 index 00000000..482ca687 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0201.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 24 * 8; + int row = vec_idx / 24; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 24 * 8; + int row_1 = vec_idx_1 / 24; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0202.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0202.cu new file mode 100644 index 00000000..7070fa90 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0202.cu @@ -0,0 +1,600 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 49152 +#define SMEM_SMEM_X_STRIDE 49152 +#define SMEM_SMEM_C_OFF 50176 +#define SMEM_SMEM_C_STAGE_BYTES 98304 +#define SMEM_SMEM_C_STRIDE 98304 +#define SMEM_SMEM_CSQ_OFF 148480 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 149504 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int smem_c_addr = smem + 50176; + float* smem_csq = reinterpret_cast(smem_raw + 148480); + const int smem_csq_addr = smem + 148480; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 49152); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 98304); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0203.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0203.cu new file mode 100644 index 00000000..482ca687 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0203.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 24 * 8; + int row = vec_idx / 24; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 24 * 8; + int row_1 = vec_idx_1 / 24; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0204.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0204.cu new file mode 100644 index 00000000..17caae28 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0204.cu @@ -0,0 +1,1211 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 49152 +#define SMEM_SMEM_X0_STRIDE 49152 +#define SMEM_SMEM_X1_OFF 50176 +#define SMEM_SMEM_X1_STAGE_BYTES 49152 +#define SMEM_SMEM_X1_STRIDE 49152 +#define SMEM_SMEM_C_OFF 99328 +#define SMEM_SMEM_C_STAGE_BYTES 98304 +#define SMEM_SMEM_C_STRIDE 98304 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int smem_x1_addr = smem + 50176; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); + const int smem_c_addr = smem + 99328; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 49152); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 49152); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 98304); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x0_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_4 = smem_x1_addr + 16384; + int _mma_a_lo_4 = make_warp_uniform((_mma_a_addr_4 >> 4) & 0x3FFF); + int _mma_b_addr_4 = smem_c_addr + 32768; + int _mma_b_lo_4 = make_warp_uniform((_mma_b_addr_4 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_5 = smem_x1_addr + 32768; + int _mma_a_lo_5 = make_warp_uniform((_mma_a_addr_5 >> 4) & 0x3FFF); + int _mma_b_addr_5 = smem_c_addr + 65536; + int _mma_b_lo_5 = make_warp_uniform((_mma_b_addr_5 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0205.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0205.cu new file mode 100644 index 00000000..482ca687 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0205.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 24 * 8; + int row = vec_idx / 24; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 24 * 8; + int row_1 = vec_idx_1 / 24; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0206.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0206.cu new file mode 100644 index 00000000..17caae28 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0206.cu @@ -0,0 +1,1211 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 49152 +#define SMEM_SMEM_X0_STRIDE 49152 +#define SMEM_SMEM_X1_OFF 50176 +#define SMEM_SMEM_X1_STAGE_BYTES 49152 +#define SMEM_SMEM_X1_STRIDE 49152 +#define SMEM_SMEM_C_OFF 99328 +#define SMEM_SMEM_C_STAGE_BYTES 98304 +#define SMEM_SMEM_C_STRIDE 98304 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int smem_x1_addr = smem + 50176; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); + const int smem_c_addr = smem + 99328; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 49152); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 49152); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 98304); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x0_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_4 = smem_x1_addr + 16384; + int _mma_a_lo_4 = make_warp_uniform((_mma_a_addr_4 >> 4) & 0x3FFF); + int _mma_b_addr_4 = smem_c_addr + 32768; + int _mma_b_lo_4 = make_warp_uniform((_mma_b_addr_4 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_5 = smem_x1_addr + 32768; + int _mma_a_lo_5 = make_warp_uniform((_mma_a_addr_5 >> 4) & 0x3FFF); + int _mma_b_addr_5 = smem_c_addr + 65536; + int _mma_b_lo_5 = make_warp_uniform((_mma_b_addr_5 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0207.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0207.cu new file mode 100644 index 00000000..482ca687 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0207.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 24 * 8; + int row = vec_idx / 24; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 24 * 8; + int row_1 = vec_idx_1 / 24; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0208.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0208.cu new file mode 100644 index 00000000..17caae28 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0208.cu @@ -0,0 +1,1211 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 49152 +#define SMEM_SMEM_X0_STRIDE 49152 +#define SMEM_SMEM_X1_OFF 50176 +#define SMEM_SMEM_X1_STAGE_BYTES 49152 +#define SMEM_SMEM_X1_STRIDE 49152 +#define SMEM_SMEM_C_OFF 99328 +#define SMEM_SMEM_C_STAGE_BYTES 98304 +#define SMEM_SMEM_C_STRIDE 98304 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int smem_x1_addr = smem + 50176; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); + const int smem_c_addr = smem + 99328; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 49152); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 49152); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 98304); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x0_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_4 = smem_x1_addr + 16384; + int _mma_a_lo_4 = make_warp_uniform((_mma_a_addr_4 >> 4) & 0x3FFF); + int _mma_b_addr_4 = smem_c_addr + 32768; + int _mma_b_lo_4 = make_warp_uniform((_mma_b_addr_4 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_5 = smem_x1_addr + 32768; + int _mma_a_lo_5 = make_warp_uniform((_mma_a_addr_5 >> 4) & 0x3FFF); + int _mma_b_addr_5 = smem_c_addr + 65536; + int _mma_b_lo_5 = make_warp_uniform((_mma_b_addr_5 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0209.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0209.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0209.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0210.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0210.cu new file mode 100644 index 00000000..bf486f4b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0210.cu @@ -0,0 +1,639 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 65536 +#define SMEM_SMEM_X_STRIDE 65536 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 131072 +#define SMEM_SMEM_C_STRIDE 131072 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 65536); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 131072); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x_addr + 49152; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 98304; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0211.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0211.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0211.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0212.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0212.cu new file mode 100644 index 00000000..bf486f4b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0212.cu @@ -0,0 +1,639 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 65536 +#define SMEM_SMEM_X_STRIDE 65536 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 131072 +#define SMEM_SMEM_C_STRIDE 131072 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 65536); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 131072); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x_addr + 49152; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 98304; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0213.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0213.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0213.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0214.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0214.cu new file mode 100644 index 00000000..bf486f4b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0214.cu @@ -0,0 +1,639 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 65536 +#define SMEM_SMEM_X_STRIDE 65536 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 131072 +#define SMEM_SMEM_C_STRIDE 131072 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 65536); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 131072); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x_addr + 49152; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 98304; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0215.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0215.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0215.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0216.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0216.cu new file mode 100644 index 00000000..bf486f4b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0216.cu @@ -0,0 +1,639 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 65536 +#define SMEM_SMEM_X_STRIDE 65536 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 131072 +#define SMEM_SMEM_C_STRIDE 131072 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 65536); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 131072); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x_addr + 49152; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 98304; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0217.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0217.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0217.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0218.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0218.cu new file mode 100644 index 00000000..bf486f4b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0218.cu @@ -0,0 +1,639 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 65536 +#define SMEM_SMEM_X_STRIDE 65536 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 131072 +#define SMEM_SMEM_C_STRIDE 131072 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 65536); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 131072); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x_addr + 49152; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 98304; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0219.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0219.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0219.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0220.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0220.cu new file mode 100644 index 00000000..bf486f4b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0220.cu @@ -0,0 +1,639 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 65536 +#define SMEM_SMEM_X_STRIDE 65536 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 131072 +#define SMEM_SMEM_C_STRIDE 131072 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 65536); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 131072); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x_addr + 49152; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 98304; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0221.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0221.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0221.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0222.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0222.cu new file mode 100644 index 00000000..bf486f4b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0222.cu @@ -0,0 +1,639 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 65536 +#define SMEM_SMEM_X_STRIDE 65536 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 131072 +#define SMEM_SMEM_C_STRIDE 131072 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 65536); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 131072); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x_addr + 49152; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 98304; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0223.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0223.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0223.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0224.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0224.cu new file mode 100644 index 00000000..bf486f4b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0224.cu @@ -0,0 +1,639 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 65536 +#define SMEM_SMEM_X_STRIDE 65536 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 131072 +#define SMEM_SMEM_C_STRIDE 131072 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 65536); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 131072); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x_addr + 49152; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 98304; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0225.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0225.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0225.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0226.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0226.cu new file mode 100644 index 00000000..bf486f4b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0226.cu @@ -0,0 +1,639 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 65536 +#define SMEM_SMEM_X_STRIDE 65536 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 131072 +#define SMEM_SMEM_C_STRIDE 131072 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 65536); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 131072); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x_addr + 49152; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 98304; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0227.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0227.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0227.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0228.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0228.cu new file mode 100644 index 00000000..bf486f4b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0228.cu @@ -0,0 +1,639 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 65536 +#define SMEM_SMEM_X_STRIDE 65536 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 131072 +#define SMEM_SMEM_C_STRIDE 131072 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 65536); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 131072); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x_addr + 49152; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 98304; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0229.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0229.cu new file mode 100644 index 00000000..fcc7623c --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0229.cu @@ -0,0 +1,520 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 4096 +#define SMEM_SMEM_X_STRIDE 4096 +#define SMEM_SMEM_C_OFF 5120 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 21504 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 22528 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 5120); + const int smem_c_addr = smem + 5120; + float* smem_csq = reinterpret_cast(smem_raw + 21504); + const int smem_csq_addr = smem + 21504; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: compute ---- + if (warp <= 3) { + { // compute_main + int num_tiles = B * num_n_tiles; + int compute_warp = warp; + int lane_pair = lane % 4; + int row_origin = compute_warp * 16; + int row_lane_base = row_origin + lane / 4; + int row0 = row_lane_base; + int row1 = row_lane_base + 8; + int compute_tid = compute_warp * 32 + lane; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 64; + int global_n0 = off_n + row0; + int global_n1 = off_n + row1; + int out_offset0 = batch * N + global_n0; + int out_offset1 = batch * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best0 = -3.4e+38f; + float best1 = -3.4e+38f; + int idx0 = 0; + int idx1 = 0; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + if ((float)compute_tid < 64.0f) { + int csq_base = compute_tid * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch * K + off_k + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + #pragma unroll + for (int score_base = 0; score_base < 256; score_base += 128) { + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, [%64];" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])), "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) + : "r"(taddr + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int rep = 0; rep < 16.0f; rep++) { + int local_reg = rep * 4; + int col_base = score_base + rep * 8 + lane_pair * 2; + float csq0 = smem_csq[col_base]; + float csq1 = smem_csq[col_base + 1]; + float d0 = _tmem_load_0[local_reg] - csq0; + if (d0 > best0) { + best0 = d0; + idx0 = off_k + col_base; + } + float d1 = _tmem_load_0[local_reg + 1] - csq1; + if (d1 > best0) { + best0 = d1; + idx0 = off_k + col_base + 1; + } + float d2 = _tmem_load_0[local_reg + 2] - csq0; + if (d2 > best1) { + best1 = d2; + idx1 = off_k + col_base; + } + float d3 = _tmem_load_0[local_reg + 3] - csq1; + if (d3 > best1) { + best1 = d3; + idx1 = off_k + col_base + 1; + } + } + } + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best0, 1); + float peer0 = _shfl_xor_0; + int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, idx0, 1); + int peer0_idx = _shfl_xor_1; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer0_idx; + } + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, best0, 2); + peer0 = _shfl_xor_2; + int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, idx0, 2); + peer0_idx = _shfl_xor_3; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer0_idx; + } + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, best1, 1); + float peer1 = _shfl_xor_4; + int _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, idx1, 1); + int peer1_idx = _shfl_xor_5; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer1_idx; + } + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, best1, 2); + peer1 = _shfl_xor_6; + int _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, idx1, 2); + peer1_idx = _shfl_xor_7; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer1_idx; + } + if (lane_pair == 0) { + if (global_n0 < N) { + *((int*)(out + out_offset0)) = idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = idx1; + } + } + } + } + // ---- Role: load ---- + } else if (warp == 4) { + { // load_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 64; + int x_row = batch_1 * N + off_n_1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + int c_row = batch_1 * K + off_k_1; + #pragma unroll 1 + for (int iter_d = 0; iter_d < 7; iter_d++) { + int d_group = iter_d; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, d_group, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 4096); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, d_group, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 5) { + { // mma_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int iter_d_1 = 0; iter_d_1 < 7; iter_d_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((iter_d_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 71304336;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0230.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0230.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0230.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0231.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0231.cu new file mode 100644 index 00000000..bf486f4b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0231.cu @@ -0,0 +1,639 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 65536 +#define SMEM_SMEM_X_STRIDE 65536 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 131072 +#define SMEM_SMEM_C_STRIDE 131072 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 65536); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 131072); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x_addr + 49152; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 98304; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0232.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0232.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0232.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0233.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0233.cu new file mode 100644 index 00000000..bf486f4b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0233.cu @@ -0,0 +1,639 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 65536 +#define SMEM_SMEM_X_STRIDE 65536 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 131072 +#define SMEM_SMEM_C_STRIDE 131072 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 65536); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 131072); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x_addr + 49152; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 98304; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0234.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0234.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0234.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0235.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0235.cu new file mode 100644 index 00000000..bf486f4b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0235.cu @@ -0,0 +1,639 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 65536 +#define SMEM_SMEM_X_STRIDE 65536 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 131072 +#define SMEM_SMEM_C_STRIDE 131072 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 65536); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 131072); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x_addr + 49152; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 98304; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0236.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0236.cu new file mode 100644 index 00000000..f2293d4d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0236.cu @@ -0,0 +1,520 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 25600 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 26624 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d288_exactd_a532_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 25600); + const int smem_csq_addr = smem + 25600; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int iter_d = 0; iter_d < 9; iter_d++) { + int d_group = iter_d; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, d_group, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, d_group, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int iter_d_1 = 0; iter_d_1 < 9; iter_d_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((iter_d_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0237.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0237.cu new file mode 100644 index 00000000..edfe179b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0237.cu @@ -0,0 +1,512 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 25600 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 26624 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, float* __restrict__ partial_scores, int* __restrict__ partial_indices, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 25600); + const int smem_csq_addr = smem + 25600; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int total_work = B * num_n_tiles * K_tiles; + int feature_tiles = 9; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_k = work_idx % (unsigned int)K_tiles; + int point_tile_idx = work_idx / (unsigned int)K_tiles; + int batch = point_tile_idx / num_n_tiles; + int n_tile = point_tile_idx % num_n_tiles; + int off_n = n_tile * 128; + int off_k = iter_k * 256; + int x_row = batch * N + off_n; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int total_work_1 = B * num_n_tiles * K_tiles; + int feature_tiles_1 = 9; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _work_idx = bid; _work_idx < total_work_1; _work_idx += num_bids) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int total_work_2 = B * num_n_tiles * K_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_2; work_idx_1 += num_bids) { + int iter_k_1 = work_idx_1 % (unsigned int)K_tiles; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_tiles; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int off_k_1 = iter_k_1 * 256; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = off_k_1; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + int partial_offset = work_idx_1 * 128 + (unsigned int)(warp % 4 * 32 + lane); + *((float*)(partial_scores + partial_offset)) = best_score; + *((int*)(partial_indices + partial_offset)) = best_idx; + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0238.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0238.cu new file mode 100644 index 00000000..0249cc88 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0238.cu @@ -0,0 +1,62 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 128 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce(float* __restrict__ partial_scores, int* __restrict__ partial_indices, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 128 + row; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int partial_offset = (point_tile_idx * (unsigned int)K_tiles + (unsigned int)iter_k) * 128 + (unsigned int)row; + float score = partial_scores[partial_offset]; + int idx = partial_indices[partial_offset]; + if (score > best_score) { + best_score = score; + best_idx = idx; + } + } + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = best_idx; + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0239.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0239.cu new file mode 100644 index 00000000..edfe179b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0239.cu @@ -0,0 +1,512 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 25600 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 26624 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, float* __restrict__ partial_scores, int* __restrict__ partial_indices, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 25600); + const int smem_csq_addr = smem + 25600; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int total_work = B * num_n_tiles * K_tiles; + int feature_tiles = 9; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_k = work_idx % (unsigned int)K_tiles; + int point_tile_idx = work_idx / (unsigned int)K_tiles; + int batch = point_tile_idx / num_n_tiles; + int n_tile = point_tile_idx % num_n_tiles; + int off_n = n_tile * 128; + int off_k = iter_k * 256; + int x_row = batch * N + off_n; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int total_work_1 = B * num_n_tiles * K_tiles; + int feature_tiles_1 = 9; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _work_idx = bid; _work_idx < total_work_1; _work_idx += num_bids) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int total_work_2 = B * num_n_tiles * K_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_2; work_idx_1 += num_bids) { + int iter_k_1 = work_idx_1 % (unsigned int)K_tiles; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_tiles; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int off_k_1 = iter_k_1 * 256; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = off_k_1; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + int partial_offset = work_idx_1 * 128 + (unsigned int)(warp % 4 * 32 + lane); + *((float*)(partial_scores + partial_offset)) = best_score; + *((int*)(partial_indices + partial_offset)) = best_idx; + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0240.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0240.cu new file mode 100644 index 00000000..0249cc88 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0240.cu @@ -0,0 +1,62 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 128 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce(float* __restrict__ partial_scores, int* __restrict__ partial_indices, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 128 + row; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int partial_offset = (point_tile_idx * (unsigned int)K_tiles + (unsigned int)iter_k) * 128 + (unsigned int)row; + float score = partial_scores[partial_offset]; + int idx = partial_indices[partial_offset]; + if (score > best_score) { + best_score = score; + best_idx = idx; + } + } + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = best_idx; + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0241.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0241.cu new file mode 100644 index 00000000..edfe179b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0241.cu @@ -0,0 +1,512 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 25600 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 26624 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, float* __restrict__ partial_scores, int* __restrict__ partial_indices, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 25600); + const int smem_csq_addr = smem + 25600; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int total_work = B * num_n_tiles * K_tiles; + int feature_tiles = 9; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_k = work_idx % (unsigned int)K_tiles; + int point_tile_idx = work_idx / (unsigned int)K_tiles; + int batch = point_tile_idx / num_n_tiles; + int n_tile = point_tile_idx % num_n_tiles; + int off_n = n_tile * 128; + int off_k = iter_k * 256; + int x_row = batch * N + off_n; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int total_work_1 = B * num_n_tiles * K_tiles; + int feature_tiles_1 = 9; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _work_idx = bid; _work_idx < total_work_1; _work_idx += num_bids) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int total_work_2 = B * num_n_tiles * K_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_2; work_idx_1 += num_bids) { + int iter_k_1 = work_idx_1 % (unsigned int)K_tiles; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_tiles; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int off_k_1 = iter_k_1 * 256; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = off_k_1; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + int partial_offset = work_idx_1 * 128 + (unsigned int)(warp % 4 * 32 + lane); + *((float*)(partial_scores + partial_offset)) = best_score; + *((int*)(partial_indices + partial_offset)) = best_idx; + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0242.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0242.cu new file mode 100644 index 00000000..0249cc88 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0242.cu @@ -0,0 +1,62 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 128 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce(float* __restrict__ partial_scores, int* __restrict__ partial_indices, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 128 + row; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int partial_offset = (point_tile_idx * (unsigned int)K_tiles + (unsigned int)iter_k) * 128 + (unsigned int)row; + float score = partial_scores[partial_offset]; + int idx = partial_indices[partial_offset]; + if (score > best_score) { + best_score = score; + best_idx = idx; + } + } + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = best_idx; + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0243.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0243.cu new file mode 100644 index 00000000..edfe179b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0243.cu @@ -0,0 +1,512 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 25600 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 26624 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, float* __restrict__ partial_scores, int* __restrict__ partial_indices, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 25600); + const int smem_csq_addr = smem + 25600; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int total_work = B * num_n_tiles * K_tiles; + int feature_tiles = 9; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_k = work_idx % (unsigned int)K_tiles; + int point_tile_idx = work_idx / (unsigned int)K_tiles; + int batch = point_tile_idx / num_n_tiles; + int n_tile = point_tile_idx % num_n_tiles; + int off_n = n_tile * 128; + int off_k = iter_k * 256; + int x_row = batch * N + off_n; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int total_work_1 = B * num_n_tiles * K_tiles; + int feature_tiles_1 = 9; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _work_idx = bid; _work_idx < total_work_1; _work_idx += num_bids) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int total_work_2 = B * num_n_tiles * K_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_2; work_idx_1 += num_bids) { + int iter_k_1 = work_idx_1 % (unsigned int)K_tiles; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_tiles; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int off_k_1 = iter_k_1 * 256; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = off_k_1; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + int partial_offset = work_idx_1 * 128 + (unsigned int)(warp % 4 * 32 + lane); + *((float*)(partial_scores + partial_offset)) = best_score; + *((int*)(partial_indices + partial_offset)) = best_idx; + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0244.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0244.cu new file mode 100644 index 00000000..0249cc88 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0244.cu @@ -0,0 +1,62 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 128 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce(float* __restrict__ partial_scores, int* __restrict__ partial_indices, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 128 + row; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int partial_offset = (point_tile_idx * (unsigned int)K_tiles + (unsigned int)iter_k) * 128 + (unsigned int)row; + float score = partial_scores[partial_offset]; + int idx = partial_indices[partial_offset]; + if (score > best_score) { + best_score = score; + best_idx = idx; + } + } + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = best_idx; + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0245.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0245.cu new file mode 100644 index 00000000..edfe179b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0245.cu @@ -0,0 +1,512 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 25600 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 26624 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, float* __restrict__ partial_scores, int* __restrict__ partial_indices, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 25600); + const int smem_csq_addr = smem + 25600; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int total_work = B * num_n_tiles * K_tiles; + int feature_tiles = 9; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_k = work_idx % (unsigned int)K_tiles; + int point_tile_idx = work_idx / (unsigned int)K_tiles; + int batch = point_tile_idx / num_n_tiles; + int n_tile = point_tile_idx % num_n_tiles; + int off_n = n_tile * 128; + int off_k = iter_k * 256; + int x_row = batch * N + off_n; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int total_work_1 = B * num_n_tiles * K_tiles; + int feature_tiles_1 = 9; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _work_idx = bid; _work_idx < total_work_1; _work_idx += num_bids) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int total_work_2 = B * num_n_tiles * K_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_2; work_idx_1 += num_bids) { + int iter_k_1 = work_idx_1 % (unsigned int)K_tiles; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_tiles; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int off_k_1 = iter_k_1 * 256; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = off_k_1; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + int partial_offset = work_idx_1 * 128 + (unsigned int)(warp % 4 * 32 + lane); + *((float*)(partial_scores + partial_offset)) = best_score; + *((int*)(partial_indices + partial_offset)) = best_idx; + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0246.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0246.cu new file mode 100644 index 00000000..0249cc88 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0246.cu @@ -0,0 +1,62 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 128 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce(float* __restrict__ partial_scores, int* __restrict__ partial_indices, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 128 + row; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int partial_offset = (point_tile_idx * (unsigned int)K_tiles + (unsigned int)iter_k) * 128 + (unsigned int)row; + float score = partial_scores[partial_offset]; + int idx = partial_indices[partial_offset]; + if (score > best_score) { + best_score = score; + best_idx = idx; + } + } + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = best_idx; + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0247.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0247.cu new file mode 100644 index 00000000..f2293d4d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0247.cu @@ -0,0 +1,520 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 25600 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 26624 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d288_exactd_a532_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 25600); + const int smem_csq_addr = smem + 25600; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int iter_d = 0; iter_d < 9; iter_d++) { + int d_group = iter_d; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, d_group, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, d_group, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int iter_d_1 = 0; iter_d_1 < 9; iter_d_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((iter_d_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0248.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0248.cu new file mode 100644 index 00000000..edfe179b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0248.cu @@ -0,0 +1,512 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 25600 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 26624 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, float* __restrict__ partial_scores, int* __restrict__ partial_indices, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 25600); + const int smem_csq_addr = smem + 25600; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int total_work = B * num_n_tiles * K_tiles; + int feature_tiles = 9; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_k = work_idx % (unsigned int)K_tiles; + int point_tile_idx = work_idx / (unsigned int)K_tiles; + int batch = point_tile_idx / num_n_tiles; + int n_tile = point_tile_idx % num_n_tiles; + int off_n = n_tile * 128; + int off_k = iter_k * 256; + int x_row = batch * N + off_n; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int total_work_1 = B * num_n_tiles * K_tiles; + int feature_tiles_1 = 9; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _work_idx = bid; _work_idx < total_work_1; _work_idx += num_bids) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int total_work_2 = B * num_n_tiles * K_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_2; work_idx_1 += num_bids) { + int iter_k_1 = work_idx_1 % (unsigned int)K_tiles; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_tiles; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int off_k_1 = iter_k_1 * 256; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = off_k_1; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + int partial_offset = work_idx_1 * 128 + (unsigned int)(warp % 4 * 32 + lane); + *((float*)(partial_scores + partial_offset)) = best_score; + *((int*)(partial_indices + partial_offset)) = best_idx; + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0249.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0249.cu new file mode 100644 index 00000000..0249cc88 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0249.cu @@ -0,0 +1,62 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 128 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce(float* __restrict__ partial_scores, int* __restrict__ partial_indices, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 128 + row; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int partial_offset = (point_tile_idx * (unsigned int)K_tiles + (unsigned int)iter_k) * 128 + (unsigned int)row; + float score = partial_scores[partial_offset]; + int idx = partial_indices[partial_offset]; + if (score > best_score) { + best_score = score; + best_idx = idx; + } + } + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = best_idx; + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0250.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0250.cu new file mode 100644 index 00000000..f2293d4d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0250.cu @@ -0,0 +1,520 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 25600 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 26624 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d288_exactd_a532_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 25600); + const int smem_csq_addr = smem + 25600; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int iter_d = 0; iter_d < 9; iter_d++) { + int d_group = iter_d; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, d_group, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, d_group, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int iter_d_1 = 0; iter_d_1 < 9; iter_d_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((iter_d_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0251.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0251.cu new file mode 100644 index 00000000..f2293d4d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0251.cu @@ -0,0 +1,520 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 25600 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 26624 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d288_exactd_a532_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 25600); + const int smem_csq_addr = smem + 25600; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int iter_d = 0; iter_d < 9; iter_d++) { + int d_group = iter_d; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, d_group, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, d_group, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int iter_d_1 = 0; iter_d_1 < 9; iter_d_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((iter_d_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0252.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0252.cu new file mode 100644 index 00000000..f2293d4d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0252.cu @@ -0,0 +1,520 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 25600 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 26624 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d288_exactd_a532_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 25600); + const int smem_csq_addr = smem + 25600; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int iter_d = 0; iter_d < 9; iter_d++) { + int d_group = iter_d; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, d_group, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, d_group, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int iter_d_1 = 0; iter_d_1 < 9; iter_d_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((iter_d_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0253.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0253.cu new file mode 100644 index 00000000..f2293d4d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0253.cu @@ -0,0 +1,520 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 25600 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 26624 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d288_exactd_a532_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 25600); + const int smem_csq_addr = smem + 25600; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int iter_d = 0; iter_d < 9; iter_d++) { + int d_group = iter_d; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, d_group, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, d_group, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int iter_d_1 = 0; iter_d_1 < 9; iter_d_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((iter_d_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0254.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0254.cu new file mode 100644 index 00000000..f2293d4d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0254.cu @@ -0,0 +1,520 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 25600 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 26624 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d288_exactd_a532_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 25600); + const int smem_csq_addr = smem + 25600; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int iter_d = 0; iter_d < 9; iter_d++) { + int d_group = iter_d; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, d_group, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, d_group, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int iter_d_1 = 0; iter_d_1 < 9; iter_d_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((iter_d_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0255.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0255.cu new file mode 100644 index 00000000..f2293d4d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0255.cu @@ -0,0 +1,520 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 25600 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 26624 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d288_exactd_a532_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 25600); + const int smem_csq_addr = smem + 25600; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int iter_d = 0; iter_d < 9; iter_d++) { + int d_group = iter_d; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, d_group, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, d_group, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int iter_d_1 = 0; iter_d_1 < 9; iter_d_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((iter_d_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0256.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0256.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0256.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0257.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0257.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0257.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0258.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0258.cu new file mode 100644 index 00000000..3f16db31 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0258.cu @@ -0,0 +1,525 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 4096 +#define SMEM_SMEM_X_STRIDE 4096 +#define SMEM_SMEM_C_OFF 5120 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 21504 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 22528 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d480_splitk_partial_d32k256_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, float* __restrict__ partial_scores, int* __restrict__ partial_indices, int B, int N, int D, int K, int num_n_tiles, int K_tiles, int K_slices) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 5120); + const int smem_c_addr = smem + 5120; + float* smem_csq = reinterpret_cast(smem_raw + 21504); + const int smem_csq_addr = smem + 21504; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: compute ---- + if (warp <= 3) { + { // compute_main + int total_work = B * num_n_tiles * K_slices; + int compute_warp = warp; + int lane_pair = lane % 4; + int row_origin = compute_warp * 16; + int row_lane_base = row_origin + lane / 4; + int row0 = row_lane_base; + int row1 = row_lane_base + 8; + int compute_tid = compute_warp * 32 + lane; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_slice = work_idx % (unsigned int)K_slices; + int point_tile_idx = work_idx / (unsigned int)K_slices; + int batch = point_tile_idx / num_n_tiles; + int slice_k_start = iter_slice; + int csq_smem_addr = smem_csq_addr; + float best0 = -3.4e+38f; + float best1 = -3.4e+38f; + int idx0 = slice_k_start * 256; + int idx1 = idx0; + #pragma unroll 1 + for (int local_k = 0; local_k < 1; local_k++) { + int iter_k = slice_k_start + local_k; + int off_k = iter_k * 256; + if ((float)compute_tid < 64.0f) { + int csq_base = compute_tid * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch * K + off_k + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + #pragma unroll + for (int score_base = 0; score_base < 256; score_base += 128) { + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, [%64];" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])), "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) + : "r"(taddr + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int rep = 0; rep < 16.0f; rep++) { + int local_reg = rep * 4; + int col_base = score_base + rep * 8 + lane_pair * 2; + float csq0 = smem_csq[col_base]; + float csq1 = smem_csq[col_base + 1]; + float d0 = _tmem_load_0[local_reg] - csq0; + if (d0 > best0) { + best0 = d0; + idx0 = off_k + col_base; + } + float d1 = _tmem_load_0[local_reg + 1] - csq1; + if (d1 > best0) { + best0 = d1; + idx0 = off_k + col_base + 1; + } + float d2 = _tmem_load_0[local_reg + 2] - csq0; + if (d2 > best1) { + best1 = d2; + idx1 = off_k + col_base; + } + float d3 = _tmem_load_0[local_reg + 3] - csq1; + if (d3 > best1) { + best1 = d3; + idx1 = off_k + col_base + 1; + } + } + } + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best0, 1); + float peer0 = _shfl_xor_0; + int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, idx0, 1); + int peer0_idx = _shfl_xor_1; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer0_idx; + } + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, best0, 2); + peer0 = _shfl_xor_2; + int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, idx0, 2); + peer0_idx = _shfl_xor_3; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer0_idx; + } + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, best1, 1); + float peer1 = _shfl_xor_4; + int _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, idx1, 1); + int peer1_idx = _shfl_xor_5; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer1_idx; + } + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, best1, 2); + peer1 = _shfl_xor_6; + int _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, idx1, 2); + peer1_idx = _shfl_xor_7; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer1_idx; + } + if (lane_pair == 0) { + int partial_offset0 = work_idx * 64 + (unsigned int)row0; + *((float*)(partial_scores + partial_offset0)) = best0; + *((int*)(partial_indices + partial_offset0)) = idx0; + int partial_offset1 = work_idx * 64 + (unsigned int)row1; + *((float*)(partial_scores + partial_offset1)) = best1; + *((int*)(partial_indices + partial_offset1)) = idx1; + } + } + } + // ---- Role: load ---- + } else if (warp == 4) { + { // load_main + int total_work_1 = B * num_n_tiles * K_slices; + int feature_tiles = D / 32; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_1; work_idx_1 += num_bids) { + int iter_slice_1 = work_idx_1 % (unsigned int)K_slices; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_slices; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int n_tile = point_tile_idx_1 % num_n_tiles; + int off_n = n_tile * 64; + int x_row = batch_1 * N + off_n; + int slice_k_start_1 = iter_slice_1; + #pragma unroll 1 + for (int local_k_1 = 0; local_k_1 < 1; local_k_1++) { + int iter_k_1 = slice_k_start_1 + local_k_1; + int off_k_1 = iter_k_1 * 256; + int c_row = batch_1 * K + off_k_1; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 4096); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 5) { + { // mma_main + int total_work_2 = B * num_n_tiles * K_slices; + int feature_tiles_1 = D / 32; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = bid; work_idx_2 < total_work_2; work_idx_2 += num_bids) { + int iter_slice_2 = work_idx_2 % (unsigned int)K_slices; + int slice_k_start_2 = iter_slice_2; + #pragma unroll 1 + for (int local_k_2 = 0; local_k_2 < 1; local_k_2++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 71304336;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0259.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0259.cu new file mode 100644 index 00000000..c242dff7 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0259.cu @@ -0,0 +1,81 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d480_splitk_reduce_d32k256_v1(float* __restrict__ partial_scores, int* __restrict__ partial_indices, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_slices) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid / 4; + int row_lane = tid % 4; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 64 + row; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_slice = row_lane; iter_slice < K_slices; iter_slice += 4) { + int partial_offset = (point_tile_idx * (unsigned int)K_slices + (unsigned int)iter_slice) * 64 + (unsigned int)row; + float score = partial_scores[partial_offset]; + int idx = partial_indices[partial_offset]; + if (score > best_score) { + best_score = score; + best_idx = idx; + } + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best_score, 1); + float peer_score = _shfl_xor_0; + int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, best_idx, 1); + int peer_idx = _shfl_xor_1; + if (peer_score > best_score) { + best_score = peer_score; + best_idx = peer_idx; + } + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, best_score, 2); + peer_score = _shfl_xor_2; + int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, best_idx, 2); + peer_idx = _shfl_xor_3; + if (peer_score > best_score) { + best_score = peer_score; + best_idx = peer_idx; + } + if (row_lane == 0) { + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = best_idx; + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0260.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0260.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0260.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0261.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0261.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0261.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0262.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0262.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0262.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0263.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0263.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0263.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0264.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0264.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0264.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0265.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0265.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0265.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0266.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0266.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0266.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0267.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0267.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0267.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0268.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0268.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0268.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0269.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0269.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0269.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0270.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0270.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0270.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0271.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0271.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0271.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0272.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0272.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0272.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0273.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0273.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0273.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0274.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0274.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0274.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0275.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0275.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0275.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0276.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0276.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0276.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0277.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0277.cu new file mode 100644 index 00000000..6d37a9aa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0277.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, float* __restrict__ partial_scores, int* __restrict__ partial_indices, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int total_work = B * num_n_tiles * K_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_k = work_idx % (unsigned int)K_tiles; + int point_tile_idx = work_idx / (unsigned int)K_tiles; + int batch = point_tile_idx / num_n_tiles; + int n_tile = point_tile_idx % num_n_tiles; + int off_n = n_tile * 128; + int off_k = iter_k * 256; + int x_row = batch * N + off_n; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int total_work_1 = B * num_n_tiles * K_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _work_idx = bid; _work_idx < total_work_1; _work_idx += num_bids) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int total_work_2 = B * num_n_tiles * K_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_2; work_idx_1 += num_bids) { + int iter_k_1 = work_idx_1 % (unsigned int)K_tiles; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_tiles; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int off_k_1 = iter_k_1 * 256; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = off_k_1; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + int partial_offset = work_idx_1 * 128 + (unsigned int)(warp % 4 * 32 + lane); + *((float*)(partial_scores + partial_offset)) = best_score; + *((int*)(partial_indices + partial_offset)) = best_idx; + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0278.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0278.cu new file mode 100644 index 00000000..c552357a --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0278.cu @@ -0,0 +1,62 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 128 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1(float* __restrict__ partial_scores, int* __restrict__ partial_indices, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 128 + row; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int partial_offset = (point_tile_idx * (unsigned int)K_tiles + (unsigned int)iter_k) * 128 + (unsigned int)row; + float score = partial_scores[partial_offset]; + int idx = partial_indices[partial_offset]; + if (score > best_score) { + best_score = score; + best_idx = idx; + } + } + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = best_idx; + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0279.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0279.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0279.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0280.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0280.cu new file mode 100644 index 00000000..6d37a9aa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0280.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, float* __restrict__ partial_scores, int* __restrict__ partial_indices, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int total_work = B * num_n_tiles * K_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_k = work_idx % (unsigned int)K_tiles; + int point_tile_idx = work_idx / (unsigned int)K_tiles; + int batch = point_tile_idx / num_n_tiles; + int n_tile = point_tile_idx % num_n_tiles; + int off_n = n_tile * 128; + int off_k = iter_k * 256; + int x_row = batch * N + off_n; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int total_work_1 = B * num_n_tiles * K_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _work_idx = bid; _work_idx < total_work_1; _work_idx += num_bids) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int total_work_2 = B * num_n_tiles * K_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_2; work_idx_1 += num_bids) { + int iter_k_1 = work_idx_1 % (unsigned int)K_tiles; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_tiles; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int off_k_1 = iter_k_1 * 256; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = off_k_1; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + int partial_offset = work_idx_1 * 128 + (unsigned int)(warp % 4 * 32 + lane); + *((float*)(partial_scores + partial_offset)) = best_score; + *((int*)(partial_indices + partial_offset)) = best_idx; + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0281.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0281.cu new file mode 100644 index 00000000..c552357a --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0281.cu @@ -0,0 +1,62 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 128 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1(float* __restrict__ partial_scores, int* __restrict__ partial_indices, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 128 + row; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int partial_offset = (point_tile_idx * (unsigned int)K_tiles + (unsigned int)iter_k) * 128 + (unsigned int)row; + float score = partial_scores[partial_offset]; + int idx = partial_indices[partial_offset]; + if (score > best_score) { + best_score = score; + best_idx = idx; + } + } + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = best_idx; + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0282.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0282.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0282.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0283.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0283.cu new file mode 100644 index 00000000..6d37a9aa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0283.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, float* __restrict__ partial_scores, int* __restrict__ partial_indices, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int total_work = B * num_n_tiles * K_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_k = work_idx % (unsigned int)K_tiles; + int point_tile_idx = work_idx / (unsigned int)K_tiles; + int batch = point_tile_idx / num_n_tiles; + int n_tile = point_tile_idx % num_n_tiles; + int off_n = n_tile * 128; + int off_k = iter_k * 256; + int x_row = batch * N + off_n; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int total_work_1 = B * num_n_tiles * K_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _work_idx = bid; _work_idx < total_work_1; _work_idx += num_bids) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int total_work_2 = B * num_n_tiles * K_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_2; work_idx_1 += num_bids) { + int iter_k_1 = work_idx_1 % (unsigned int)K_tiles; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_tiles; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int off_k_1 = iter_k_1 * 256; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = off_k_1; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + int partial_offset = work_idx_1 * 128 + (unsigned int)(warp % 4 * 32 + lane); + *((float*)(partial_scores + partial_offset)) = best_score; + *((int*)(partial_indices + partial_offset)) = best_idx; + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0284.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0284.cu new file mode 100644 index 00000000..c552357a --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0284.cu @@ -0,0 +1,62 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 128 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1(float* __restrict__ partial_scores, int* __restrict__ partial_indices, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 128 + row; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int partial_offset = (point_tile_idx * (unsigned int)K_tiles + (unsigned int)iter_k) * 128 + (unsigned int)row; + float score = partial_scores[partial_offset]; + int idx = partial_indices[partial_offset]; + if (score > best_score) { + best_score = score; + best_idx = idx; + } + } + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = best_idx; + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0285.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0285.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0285.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0286.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0286.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0286.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0287.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0287.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0287.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0288.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0288.cu new file mode 100644 index 00000000..6d37a9aa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0288.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, float* __restrict__ partial_scores, int* __restrict__ partial_indices, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int total_work = B * num_n_tiles * K_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_k = work_idx % (unsigned int)K_tiles; + int point_tile_idx = work_idx / (unsigned int)K_tiles; + int batch = point_tile_idx / num_n_tiles; + int n_tile = point_tile_idx % num_n_tiles; + int off_n = n_tile * 128; + int off_k = iter_k * 256; + int x_row = batch * N + off_n; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int total_work_1 = B * num_n_tiles * K_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _work_idx = bid; _work_idx < total_work_1; _work_idx += num_bids) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int total_work_2 = B * num_n_tiles * K_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_2; work_idx_1 += num_bids) { + int iter_k_1 = work_idx_1 % (unsigned int)K_tiles; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_tiles; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int off_k_1 = iter_k_1 * 256; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = off_k_1; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + int partial_offset = work_idx_1 * 128 + (unsigned int)(warp % 4 * 32 + lane); + *((float*)(partial_scores + partial_offset)) = best_score; + *((int*)(partial_indices + partial_offset)) = best_idx; + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0289.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0289.cu new file mode 100644 index 00000000..c552357a --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0289.cu @@ -0,0 +1,62 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 128 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1(float* __restrict__ partial_scores, int* __restrict__ partial_indices, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 128 + row; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int partial_offset = (point_tile_idx * (unsigned int)K_tiles + (unsigned int)iter_k) * 128 + (unsigned int)row; + float score = partial_scores[partial_offset]; + int idx = partial_indices[partial_offset]; + if (score > best_score) { + best_score = score; + best_idx = idx; + } + } + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = best_idx; + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0290.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0290.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0290.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0291.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0291.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0291.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0292.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0292.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0292.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0293.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0293.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0293.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0294.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0294.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0294.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0295.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0295.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0295.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0296.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0296.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0296.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0297.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0297.cu new file mode 100644 index 00000000..6d37a9aa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0297.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, float* __restrict__ partial_scores, int* __restrict__ partial_indices, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int total_work = B * num_n_tiles * K_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_k = work_idx % (unsigned int)K_tiles; + int point_tile_idx = work_idx / (unsigned int)K_tiles; + int batch = point_tile_idx / num_n_tiles; + int n_tile = point_tile_idx % num_n_tiles; + int off_n = n_tile * 128; + int off_k = iter_k * 256; + int x_row = batch * N + off_n; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int total_work_1 = B * num_n_tiles * K_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _work_idx = bid; _work_idx < total_work_1; _work_idx += num_bids) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int total_work_2 = B * num_n_tiles * K_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_2; work_idx_1 += num_bids) { + int iter_k_1 = work_idx_1 % (unsigned int)K_tiles; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_tiles; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int off_k_1 = iter_k_1 * 256; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = off_k_1; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + int partial_offset = work_idx_1 * 128 + (unsigned int)(warp % 4 * 32 + lane); + *((float*)(partial_scores + partial_offset)) = best_score; + *((int*)(partial_indices + partial_offset)) = best_idx; + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0298.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0298.cu new file mode 100644 index 00000000..c552357a --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0298.cu @@ -0,0 +1,62 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 128 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1(float* __restrict__ partial_scores, int* __restrict__ partial_indices, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 128 + row; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int partial_offset = (point_tile_idx * (unsigned int)K_tiles + (unsigned int)iter_k) * 128 + (unsigned int)row; + float score = partial_scores[partial_offset]; + int idx = partial_indices[partial_offset]; + if (score > best_score) { + best_score = score; + best_idx = idx; + } + } + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = best_idx; + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0299.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0299.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0299.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0300.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0300.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0300.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0301.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0301.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0301.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0302.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0302.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0302.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0303.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0303.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0303.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0304.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0304.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0304.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0305.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0305.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0305.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0306.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0306.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0306.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0307.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0307.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0307.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0308.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0308.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0308.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0309.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0309.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0309.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0310.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0310.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0310.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0311.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0311.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0311.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0312.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0312.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0312.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0313.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0313.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0313.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0314.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0314.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0314.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0315.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0315.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0315.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0316.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0316.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0316.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0317.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0317.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0317.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0318.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0318.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0318.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0319.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0319.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0319.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0320.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0320.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0320.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0321.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0321.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0321.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0322.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0322.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0322.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0323.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0323.cu new file mode 100644 index 00000000..f6311654 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0323.cu @@ -0,0 +1,535 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 41984 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 43008 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitk_partial_blockn64_g2r4_b5a6_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, float* __restrict__ partial_scores, int* __restrict__ partial_indices, int B, int N, int D, int K, int num_n_tiles, int K_tiles, int K_slices) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 41984); + const int smem_csq_addr = smem + 41984; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: compute ---- + if (warp <= 3) { + { // compute_main + int total_work = B * num_n_tiles * K_slices; + int compute_warp = warp; + int lane_pair = lane % 4; + int row_origin = compute_warp * 16; + int row_lane_base = row_origin + lane / 4; + int row0 = row_lane_base; + int row1 = row_lane_base + 8; + int compute_tid = compute_warp * 32 + lane; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_slice = work_idx % (unsigned int)K_slices; + int point_tile_idx = work_idx / (unsigned int)K_slices; + int batch = point_tile_idx / num_n_tiles; + int slice_k_start = iter_slice * 2; + int csq_smem_addr = smem_csq_addr; + float best0 = -3.4e+38f; + float best1 = -3.4e+38f; + int idx0 = slice_k_start * 256; + int idx1 = idx0; + #pragma unroll 1 + for (int local_k = 0; local_k < 2; local_k++) { + int iter_k = slice_k_start + local_k; + int off_k = iter_k * 256; + if ((float)compute_tid < 64.0f) { + int csq_base = compute_tid * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch * K + off_k + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + #pragma unroll + for (int score_base = 0; score_base < 256; score_base += 128) { + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, [%64];" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])), "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) + : "r"(taddr + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int rep = 0; rep < 16.0f; rep++) { + int local_reg = rep * 4; + int col_base = score_base + rep * 8 + lane_pair * 2; + float csq0 = smem_csq[col_base]; + float csq1 = smem_csq[col_base + 1]; + float d0 = _tmem_load_0[local_reg] - csq0; + if (d0 > best0) { + best0 = d0; + idx0 = off_k + col_base; + } + float d1 = _tmem_load_0[local_reg + 1] - csq1; + if (d1 > best0) { + best0 = d1; + idx0 = off_k + col_base + 1; + } + float d2 = _tmem_load_0[local_reg + 2] - csq0; + if (d2 > best1) { + best1 = d2; + idx1 = off_k + col_base; + } + float d3 = _tmem_load_0[local_reg + 3] - csq1; + if (d3 > best1) { + best1 = d3; + idx1 = off_k + col_base + 1; + } + } + } + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best0, 1); + float peer0 = _shfl_xor_0; + int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, idx0, 1); + int peer0_idx = _shfl_xor_1; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer0_idx; + } + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, best0, 2); + peer0 = _shfl_xor_2; + int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, idx0, 2); + peer0_idx = _shfl_xor_3; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer0_idx; + } + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, best1, 1); + float peer1 = _shfl_xor_4; + int _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, idx1, 1); + int peer1_idx = _shfl_xor_5; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer1_idx; + } + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, best1, 2); + peer1 = _shfl_xor_6; + int _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, idx1, 2); + peer1_idx = _shfl_xor_7; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer1_idx; + } + if (lane_pair == 0) { + int partial_offset0 = work_idx * 64 + (unsigned int)row0; + *((float*)(partial_scores + partial_offset0)) = best0; + *((int*)(partial_indices + partial_offset0)) = idx0; + int partial_offset1 = work_idx * 64 + (unsigned int)row1; + *((float*)(partial_scores + partial_offset1)) = best1; + *((int*)(partial_indices + partial_offset1)) = idx1; + } + } + } + // ---- Role: load ---- + } else if (warp == 4) { + { // load_main + int total_work_1 = B * num_n_tiles * K_slices; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_1; work_idx_1 += num_bids) { + int iter_slice_1 = work_idx_1 % (unsigned int)K_slices; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_slices; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int n_tile = point_tile_idx_1 % num_n_tiles; + int off_n = n_tile * 64; + int x_row = batch_1 * N + off_n; + int slice_k_start_1 = iter_slice_1 * 2; + #pragma unroll 1 + for (int local_k_1 = 0; local_k_1 < 2; local_k_1++) { + int iter_k_1 = slice_k_start_1 + local_k_1; + int off_k_1 = iter_k_1 * 256; + int c_row = batch_1 * K + off_k_1; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 5) { + { // mma_main + int total_work_2 = B * num_n_tiles * K_slices; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = bid; work_idx_2 < total_work_2; work_idx_2 += num_bids) { + int iter_slice_2 = work_idx_2 % (unsigned int)K_slices; + int slice_k_start_2 = iter_slice_2 * 2; + #pragma unroll 1 + for (int local_k_2 = 0; local_k_2 < 2; local_k_2++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 71304336;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0324.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0324.cu new file mode 100644 index 00000000..461424c3 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0324.cu @@ -0,0 +1,81 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_highd_splitk_reduce_blockn64_g2r4_b5a6_v1(float* __restrict__ partial_scores, int* __restrict__ partial_indices, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_slices) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid / 4; + int row_lane = tid % 4; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 64 + row; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_slice = row_lane; iter_slice < K_slices; iter_slice += 4) { + int partial_offset = (point_tile_idx * (unsigned int)K_slices + (unsigned int)iter_slice) * 64 + (unsigned int)row; + float score = partial_scores[partial_offset]; + int idx = partial_indices[partial_offset]; + if (score > best_score) { + best_score = score; + best_idx = idx; + } + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best_score, 1); + float peer_score = _shfl_xor_0; + int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, best_idx, 1); + int peer_idx = _shfl_xor_1; + if (peer_score > best_score) { + best_score = peer_score; + best_idx = peer_idx; + } + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, best_score, 2); + peer_score = _shfl_xor_2; + int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, best_idx, 2); + peer_idx = _shfl_xor_3; + if (peer_score > best_score) { + best_score = peer_score; + best_idx = peer_idx; + } + if (row_lane == 0) { + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = best_idx; + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0325.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0325.cu new file mode 100644 index 00000000..6d37a9aa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0325.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, float* __restrict__ partial_scores, int* __restrict__ partial_indices, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int total_work = B * num_n_tiles * K_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_k = work_idx % (unsigned int)K_tiles; + int point_tile_idx = work_idx / (unsigned int)K_tiles; + int batch = point_tile_idx / num_n_tiles; + int n_tile = point_tile_idx % num_n_tiles; + int off_n = n_tile * 128; + int off_k = iter_k * 256; + int x_row = batch * N + off_n; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int total_work_1 = B * num_n_tiles * K_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _work_idx = bid; _work_idx < total_work_1; _work_idx += num_bids) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int total_work_2 = B * num_n_tiles * K_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_2; work_idx_1 += num_bids) { + int iter_k_1 = work_idx_1 % (unsigned int)K_tiles; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_tiles; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int off_k_1 = iter_k_1 * 256; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = off_k_1; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + int partial_offset = work_idx_1 * 128 + (unsigned int)(warp % 4 * 32 + lane); + *((float*)(partial_scores + partial_offset)) = best_score; + *((int*)(partial_indices + partial_offset)) = best_idx; + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0326.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0326.cu new file mode 100644 index 00000000..c552357a --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0326.cu @@ -0,0 +1,62 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 128 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1(float* __restrict__ partial_scores, int* __restrict__ partial_indices, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 128 + row; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int partial_offset = (point_tile_idx * (unsigned int)K_tiles + (unsigned int)iter_k) * 128 + (unsigned int)row; + float score = partial_scores[partial_offset]; + int idx = partial_indices[partial_offset]; + if (score > best_score) { + best_score = score; + best_idx = idx; + } + } + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = best_idx; + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0327.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0327.cu new file mode 100644 index 00000000..3c2f58da --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0327.cu @@ -0,0 +1,531 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitd_6fcf_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0328.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0328.cu new file mode 100644 index 00000000..fcaac275 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0328.cu @@ -0,0 +1,535 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 41984 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 43008 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_highd_splitk_partial_blockn64_g1r4_streamdep_r63_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, float* __restrict__ partial_scores, int* __restrict__ partial_indices, int B, int N, int D, int K, int num_n_tiles, int K_tiles, int K_slices) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 41984); + const int smem_csq_addr = smem + 41984; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: compute ---- + if (warp <= 3) { + { // compute_main + int total_work = B * num_n_tiles * K_slices; + int compute_warp = warp; + int lane_pair = lane % 4; + int row_origin = compute_warp * 16; + int row_lane_base = row_origin + lane / 4; + int row0 = row_lane_base; + int row1 = row_lane_base + 8; + int compute_tid = compute_warp * 32 + lane; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_slice = work_idx % (unsigned int)K_slices; + int point_tile_idx = work_idx / (unsigned int)K_slices; + int batch = point_tile_idx / num_n_tiles; + int slice_k_start = iter_slice; + int csq_smem_addr = smem_csq_addr; + float best0 = -3.4e+38f; + float best1 = -3.4e+38f; + int idx0 = slice_k_start * 256; + int idx1 = idx0; + #pragma unroll 1 + for (int local_k = 0; local_k < 1; local_k++) { + int iter_k = slice_k_start + local_k; + int off_k = iter_k * 256; + if ((float)compute_tid < 64.0f) { + int csq_base = compute_tid * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch * K + off_k + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + #pragma unroll + for (int score_base = 0; score_base < 256; score_base += 128) { + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, [%64];" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])), "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) + : "r"(taddr + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int rep = 0; rep < 16.0f; rep++) { + int local_reg = rep * 4; + int col_base = score_base + rep * 8 + lane_pair * 2; + float csq0 = smem_csq[col_base]; + float csq1 = smem_csq[col_base + 1]; + float d0 = _tmem_load_0[local_reg] - csq0; + if (d0 > best0) { + best0 = d0; + idx0 = off_k + col_base; + } + float d1 = _tmem_load_0[local_reg + 1] - csq1; + if (d1 > best0) { + best0 = d1; + idx0 = off_k + col_base + 1; + } + float d2 = _tmem_load_0[local_reg + 2] - csq0; + if (d2 > best1) { + best1 = d2; + idx1 = off_k + col_base; + } + float d3 = _tmem_load_0[local_reg + 3] - csq1; + if (d3 > best1) { + best1 = d3; + idx1 = off_k + col_base + 1; + } + } + } + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best0, 1); + float peer0 = _shfl_xor_0; + int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, idx0, 1); + int peer0_idx = _shfl_xor_1; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer0_idx; + } + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, best0, 2); + peer0 = _shfl_xor_2; + int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, idx0, 2); + peer0_idx = _shfl_xor_3; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer0_idx; + } + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, best1, 1); + float peer1 = _shfl_xor_4; + int _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, idx1, 1); + int peer1_idx = _shfl_xor_5; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer1_idx; + } + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, best1, 2); + peer1 = _shfl_xor_6; + int _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, idx1, 2); + peer1_idx = _shfl_xor_7; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer1_idx; + } + if (lane_pair == 0) { + int partial_offset0 = work_idx * 64 + (unsigned int)row0; + *((float*)(partial_scores + partial_offset0)) = best0; + *((int*)(partial_indices + partial_offset0)) = idx0; + int partial_offset1 = work_idx * 64 + (unsigned int)row1; + *((float*)(partial_scores + partial_offset1)) = best1; + *((int*)(partial_indices + partial_offset1)) = idx1; + } + } + } + // ---- Role: load ---- + } else if (warp == 4) { + { // load_main + int total_work_1 = B * num_n_tiles * K_slices; + int feature_tiles = D / 64; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_1; work_idx_1 += num_bids) { + int iter_slice_1 = work_idx_1 % (unsigned int)K_slices; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_slices; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int n_tile = point_tile_idx_1 % num_n_tiles; + int off_n = n_tile * 64; + int x_row = batch_1 * N + off_n; + int slice_k_start_1 = iter_slice_1; + #pragma unroll 1 + for (int local_k_1 = 0; local_k_1 < 1; local_k_1++) { + int iter_k_1 = slice_k_start_1 + local_k_1; + int off_k_1 = iter_k_1 * 256; + int c_row = batch_1 * K + off_k_1; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 5) { + { // mma_main + int total_work_2 = B * num_n_tiles * K_slices; + int feature_tiles_1 = D / 64; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = bid; work_idx_2 < total_work_2; work_idx_2 += num_bids) { + int iter_slice_2 = work_idx_2 % (unsigned int)K_slices; + int slice_k_start_2 = iter_slice_2; + #pragma unroll 1 + for (int local_k_2 = 0; local_k_2 < 1; local_k_2++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 71304336;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0329.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0329.cu new file mode 100644 index 00000000..254d8e9d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0329.cu @@ -0,0 +1,81 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_highd_splitk_reduce_blockn64_g1r4_streamdep_r63_v1(float* __restrict__ partial_scores, int* __restrict__ partial_indices, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_slices) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid / 4; + int row_lane = tid % 4; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 64 + row; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_slice = row_lane; iter_slice < K_slices; iter_slice += 4) { + int partial_offset = (point_tile_idx * (unsigned int)K_slices + (unsigned int)iter_slice) * 64 + (unsigned int)row; + float score = partial_scores[partial_offset]; + int idx = partial_indices[partial_offset]; + if (score > best_score) { + best_score = score; + best_idx = idx; + } + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best_score, 1); + float peer_score = _shfl_xor_0; + int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, best_idx, 1); + int peer_idx = _shfl_xor_1; + if (peer_score > best_score) { + best_score = peer_score; + best_idx = peer_idx; + } + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, best_score, 2); + peer_score = _shfl_xor_2; + int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, best_idx, 2); + peer_idx = _shfl_xor_3; + if (peer_score > best_score) { + best_score = peer_score; + best_idx = peer_idx; + } + if (row_lane == 0) { + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = best_idx; + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0330.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0330.cu new file mode 100644 index 00000000..3e4e2222 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0330.cu @@ -0,0 +1,483 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORES_OFFSET 0 +#define NUM_X_PIPE_STAGES 3 +#define NUM_C_PIPE_STAGES 4 +#define NUM_SCORE_PIPE_STAGES 4 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 4096 +#define SMEM_SMEM_X_STRIDE 4096 +#define SMEM_SMEM_C_OFF 13312 +#define SMEM_SMEM_C_STAGE_BYTES 2048 +#define SMEM_SMEM_C_STRIDE 2048 +#define SMEM_SMEM_CSQ_OFF 21504 +#define SMEM_SMEM_CSQ_STAGE_BYTES 2048 +#define SMEM_SMEM_CSQ_STRIDE 2048 +#define SMEM_TOTAL 23552 +#define THREADS 192 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_microdim_d16_pipeline4_08f9_v4(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 13312); + const int smem_c_addr = smem + 13312; + float* smem_csq = reinterpret_cast(smem_raw + 21504); + const int smem_csq_addr = smem + 21504; + + // Mbarrier init (6 groups, 22 barriers) + // Mbarriers at smem_raw[0..176) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // --- pipeline 'x_pipe' --- + // x_full: 3 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + mbarrier_init_pred(smem + 8, 1, leader); + mbarrier_init_pred(smem + 16, 1, leader); + // x_empty: 3 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + mbarrier_init_pred(smem + 32, 1, leader); + mbarrier_init_pred(smem + 40, 1, leader); + // --- pipeline 'c_pipe' --- + // c_full: 4 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + mbarrier_init_pred(smem + 56, 1, leader); + mbarrier_init_pred(smem + 64, 1, leader); + mbarrier_init_pred(smem + 72, 1, leader); + // c_empty: 4 barriers, init_count=1 + mbarrier_init_pred(smem + 80, 1, leader); + mbarrier_init_pred(smem + 88, 1, leader); + mbarrier_init_pred(smem + 96, 1, leader); + mbarrier_init_pred(smem + 104, 1, leader); + // --- pipeline 'score_pipe' --- + // score_full: 4 barriers, init_count=1 + mbarrier_init_pred(smem + 112, 1, leader); + mbarrier_init_pred(smem + 120, 1, leader); + mbarrier_init_pred(smem + 128, 1, leader); + mbarrier_init_pred(smem + 136, 1, leader); + // score_empty: 4 barriers, init_count=4 + mbarrier_init_pred(smem + 144, 4, leader); + mbarrier_init_pred(smem + 152, 4, leader); + mbarrier_init_pred(smem + 160, 4, leader); + mbarrier_init_pred(smem + 168, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 176); + if (warp == 0) { + int _tmem_hold = smem + 176; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 48) + #define c_empty_addr (mbar_base + 80) + #define score_full_addr (mbar_base + 112) + #define score_empty_addr (mbar_base + 144) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_scores = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + unsigned int x_stage = 0; + unsigned int c_stage = 0; + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty = 1; + unsigned int _phase_c_empty = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int x_row = batch * N + n_tile * 128; + mbarrier_wait(x_empty_addr + (x_stage) * 8, _phase_x_empty); + tma_3d_gmem2smem(smem_x_addr + x_stage * 4096, x_tmap, 0, x_row, 0, x_full_addr + (x_stage) * 8); + mbarrier_arrive_expect_tx(x_full_addr + (x_stage) * 8, 4096); + x_stage = (x_stage + 1) % 3; + if (x_stage == 0) { _phase_x_empty ^= 1; } + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int c_row = batch * K + iter_k * 64; + mbarrier_wait(c_empty_addr + (c_stage) * 8, _phase_c_empty); + tma_3d_gmem2smem(smem_c_addr + c_stage * 2048, c_tmap, 0, c_row, 0, c_full_addr + (c_stage) * 8); + mbarrier_arrive_expect_tx(c_full_addr + (c_stage) * 8, 2048); + c_stage = (c_stage + 1) % 4; + if (c_stage == 0) { _phase_c_empty ^= 1; } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + unsigned int x_stage_1 = 0; + unsigned int c_stage_1 = 0; + unsigned int score_stage = 0; + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full = 0; + unsigned int _phase_c_full = 0; + unsigned int _phase_score_empty = 1; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr + (x_stage_1) * 8, _phase_x_full); + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(c_full_addr + (c_stage_1) * 8, _phase_c_full); + mbarrier_wait(score_empty_addr + (score_stage) * 8, _phase_score_empty); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr + x_stage_1 * 4096; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr + c_stage_1 * 2048; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + mma_ss_step(_mma_a_lo_0, _mma_b_lo_0, (tmem_scores + (score_stage * 64)), 135267472, 0, 0xC0004010U, 0xC0004010U); + elect_commit(score_full_addr + (score_stage) * 8); + elect_commit(c_empty_addr + (c_stage_1) * 8); + c_stage_1 = (c_stage_1 + 1) % 4; + if (c_stage_1 == 0) { _phase_c_full ^= 1; } + score_stage = (score_stage + 1) % 4; + if (score_stage == 0) { _phase_score_empty ^= 1; } + } + if (elect_sync()) { + mbarrier_arrive(x_empty_addr + (x_stage_1) * 8); + } + x_stage_1 = (x_stage_1 + 1) % 3; + if (x_stage_1 == 0) { _phase_x_full ^= 1; } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + unsigned int score_stage_1 = 0; + int num_tiles_2 = B * num_n_tiles; + int row = warp % 4 * 32 + lane; + int row_base = warp % 4 * 32 << 16; + unsigned int _phase_score_full = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int out_offset = batch_1 * N + n_tile_1 * 128 + row; + float best0 = -3.4e+38f; + float best1 = -3.4e+38f; + float best2 = -3.4e+38f; + float best3 = -3.4e+38f; + int idx0 = 0; + int idx1 = 0; + int idx2 = 0; + int idx3 = 0; + int csq_base = row * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(smem_csq_addr + (unsigned int)(csq_base * 4)), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int k_base = iter_k_1 * 64; + mbarrier_wait(score_full_addr + (score_stage_1) * 8, _phase_score_full); + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x64.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, [%64];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]) + : "r"(taddr + (unsigned int)row_base + score_stage_1 * 64) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int j_base = 0; j_base < 64; j_base += 4) { + float d0 = _tmem_load_0[j_base] - smem_csq[k_base + j_base]; + if (d0 > best0) { + best0 = d0; + idx0 = k_base + j_base; + } + float d1 = _tmem_load_0[j_base + 1] - smem_csq[k_base + j_base + 1]; + if (d1 > best1) { + best1 = d1; + idx1 = k_base + j_base + 1; + } + float d2 = _tmem_load_0[j_base + 2] - smem_csq[k_base + j_base + 2]; + if (d2 > best2) { + best2 = d2; + idx2 = k_base + j_base + 2; + } + float d3 = _tmem_load_0[j_base + 3] - smem_csq[k_base + j_base + 3]; + if (d3 > best3) { + best3 = d3; + idx3 = k_base + j_base + 3; + } + } + if (elect_sync()) { + mbarrier_arrive(score_empty_addr + (score_stage_1) * 8); + } + score_stage_1 = (score_stage_1 + 1) % 4; + if (score_stage_1 == 0) { _phase_score_full ^= 1; } + } + if (best1 > best0) { + best0 = best1; + idx0 = idx1; + } + if (best3 > best2) { + best2 = best3; + idx2 = idx3; + } + if (best2 > best0) { + best0 = best2; + idx0 = idx2; + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + *((int*)(out + out_offset)) = idx0; + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0331.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0331.cu new file mode 100644 index 00000000..3778419b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0331.cu @@ -0,0 +1,520 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 2048 +#define SMEM_SMEM_CSQ_STRIDE 2048 +#define SMEM_TOTAL 52224 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_microdim_raw_tma_08f9_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_2d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_2d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (off_k_1 + score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (off_k_1 + score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0332.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0332.cu new file mode 100644 index 00000000..fe28d85a --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0332.cu @@ -0,0 +1,577 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_microdim_direct_9c0d_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + int lane_start = tid; + const int x_vecs = 1024; + const int c_vecs = 2048; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row_base = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + #pragma unroll 1 + for (unsigned int vec_idx = lane_start; vec_idx < x_vecs; vec_idx += 32) { + int d_pad = vec_idx % 8 * 8; + int row = vec_idx / 8; + float vals[8]; + unsigned int packed[4]; + if (d_pad < D) { + { + const uint4* _vptr_0 = reinterpret_cast(x + ((x_row_base + row) * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + vals[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + } else { + #pragma unroll + for (int vi = 0; vi < 8; vi++) { + vals[vi] = 0.0f; + } + } + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(vals[_lp*2 + 0], vals[_lp*2+1 + 0])); + packed[_lp] = *(uint32_t*)&_bf2; + } + int x_addr = (smem_x_addr + (unsigned int)(d_pad / 64 * 16384 + row * 128 + d_pad % 64 * 2 ^ (d_pad / 64 * 16384 + row * 128 + d_pad % 64 * 2 >> 7 & 7) << 4)); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(x_addr), "r"(packed[0]), "r"(packed[1]), "r"(packed[2]), "r"(packed[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(32)); + if (warp == 0) { + if (elect_sync()) { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(x_full_addr); + } + } + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row_base = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + #pragma unroll 1 + for (unsigned int vec_idx_1 = lane_start; vec_idx_1 < c_vecs; vec_idx_1 += 32) { + int d_pad_1 = vec_idx_1 % 8 * 8; + int row_1 = vec_idx_1 / 8; + float vals_1[8]; + unsigned int packed_1[4]; + if (d_pad_1 < D) { + { + const uint4* _vptr_1 = reinterpret_cast(centroids + ((c_row_base + row_1) * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + vals_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + } else { + #pragma unroll + for (int vi_1 = 0; vi_1 < 8; vi_1++) { + vals_1[vi_1] = 0.0f; + } + } + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(vals_1[_lp*2 + 0], vals_1[_lp*2+1 + 0])); + packed_1[_lp] = *(uint32_t*)&_bf2; + } + int c_addr = (smem_c_addr + (unsigned int)(d_pad_1 / 64 * 32768 + row_1 * 128 + d_pad_1 % 64 * 2 ^ (d_pad_1 / 64 * 32768 + row_1 * 128 + d_pad_1 % 64 * 2 >> 7 & 7) << 4)); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(c_addr), "r"(packed_1[0]), "r"(packed_1[1]), "r"(packed_1[2]), "r"(packed_1[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(32)); + if (warp == 0) { + if (elect_sync()) { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(c_full_addr); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 10, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0333.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0333.cu new file mode 100644 index 00000000..5b836885 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0333.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_microdim_pack_6cd2_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 8 * 8; + int row = vec_idx / 8; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 8 * 8; + int row_1 = vec_idx_1 / 8; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0334.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0334.cu new file mode 100644 index 00000000..1c18f40d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0334.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_microdim_6cd2_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0335.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0335.cu new file mode 100644 index 00000000..df9612f5 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0335.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_lowdim_pack_e50c_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 16 * 8; + int row = vec_idx / 16; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 16 * 8; + int row_1 = vec_idx_1 / 16; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0336.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0336.cu new file mode 100644 index 00000000..1c18f40d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0336.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_microdim_6cd2_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0337.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0337.cu new file mode 100644 index 00000000..bd103c07 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0337.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_lowdim_e50c_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0338.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0338.cu new file mode 100644 index 00000000..b0ab1f24 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0338.cu @@ -0,0 +1,139 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_gap_pad_pack_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int D_PAD, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int d_pad_vecs = D_PAD / 8; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % (unsigned int)d_pad_vecs * 8; + int row = vec_idx / (unsigned int)d_pad_vecs; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % (unsigned int)d_pad_vecs * 8; + int row_1 = vec_idx_1 / (unsigned int)d_pad_vecs; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0339.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0339.cu new file mode 100644 index 00000000..a6df74d2 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0339.cu @@ -0,0 +1,522 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 16384 +#define SMEM_SMEM_X_STRIDE 16384 +#define SMEM_SMEM_C_OFF 17408 +#define SMEM_SMEM_C_STAGE_BYTES 32768 +#define SMEM_SMEM_C_STRIDE 32768 +#define SMEM_SMEM_CSQ_OFF 50176 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 51200 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_c_addr = smem + 17408; + float* smem_csq = reinterpret_cast(smem_raw + 50176); + const int smem_csq_addr = smem + 50176; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 16384); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 32768); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _tile_idx = bid; _tile_idx < num_tiles_1; _tile_idx += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int _iter_k = 0; _iter_k < K_tiles; _iter_k++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_2; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0340.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0340.cu new file mode 100644 index 00000000..d5f77bfa --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0340.cu @@ -0,0 +1,1133 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 32768 +#define SMEM_SMEM_X0_STRIDE 32768 +#define SMEM_SMEM_X1_OFF 33792 +#define SMEM_SMEM_X1_STAGE_BYTES 32768 +#define SMEM_SMEM_X1_STRIDE 32768 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 132096 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 133120 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v15(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_x1_addr = smem + 33792; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 132096); + const int smem_csq_addr = smem + 132096; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 32768); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x1_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr + 16384; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 32768; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0341.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0341.cu new file mode 100644 index 00000000..de141964 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0341.cu @@ -0,0 +1,561 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 32768 +#define SMEM_SMEM_X_STRIDE 32768 +#define SMEM_SMEM_C_OFF 33792 +#define SMEM_SMEM_C_STAGE_BYTES 65536 +#define SMEM_SMEM_C_STRIDE 65536 +#define SMEM_SMEM_CSQ_OFF 99328 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 100352 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_v10(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_c_addr = smem + 33792; + float* smem_csq = reinterpret_cast(smem_raw + 99328); + const int smem_csq_addr = smem + 99328; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 32768); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 65536); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0342.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0342.cu new file mode 100644 index 00000000..bf486f4b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0342.cu @@ -0,0 +1,639 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 65536 +#define SMEM_SMEM_X_STRIDE 65536 +#define SMEM_SMEM_C_OFF 66560 +#define SMEM_SMEM_C_STAGE_BYTES 131072 +#define SMEM_SMEM_C_STRIDE 131072 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_c_addr = smem + 66560; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 65536); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 131072); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x_addr + 49152; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr + 98304; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0343.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0343.cu new file mode 100644 index 00000000..17caae28 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0343.cu @@ -0,0 +1,1211 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 512 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X0_OFF 1024 +#define SMEM_SMEM_X0_STAGE_BYTES 49152 +#define SMEM_SMEM_X0_STRIDE 49152 +#define SMEM_SMEM_X1_OFF 50176 +#define SMEM_SMEM_X1_STAGE_BYTES 49152 +#define SMEM_SMEM_X1_STRIDE 49152 +#define SMEM_SMEM_C_OFF 99328 +#define SMEM_SMEM_C_STAGE_BYTES 98304 +#define SMEM_SMEM_C_STRIDE 98304 +#define SMEM_SMEM_CSQ_OFF 197632 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 198656 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x0_addr = smem + 1024; + __nv_bfloat16* smem_x1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int smem_x1_addr = smem + 50176; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); + const int smem_c_addr = smem + 99328; + float* smem_csq = reinterpret_cast(smem_raw + 197632); + const int smem_csq_addr = smem + 197632; + + // Mbarrier init (10 groups, 10 barriers) + // Mbarriers at smem_raw[0..80) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x0_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // x1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // x1_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + // score0_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 48, 1, leader); + // score0_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 56, 4, leader); + // score1_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // score1_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 72, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 80); + if (warp == 0) { + int _tmem_hold = smem + 80; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x0_full_addr (mbar_base + 0) + #define x0_empty_addr (mbar_base + 8) + #define x1_full_addr (mbar_base + 16) + #define x1_empty_addr (mbar_base + 24) + #define c_full_addr (mbar_base + 32) + #define c_empty_addr (mbar_base + 40) + #define score0_full_addr (mbar_base + 48) + #define score0_empty_addr (mbar_base + 56) + #define score1_full_addr (mbar_base + 64) + #define score1_empty_addr (mbar_base + 72) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int pair_n_tiles = num_n_tiles / 2; + int num_tiles = B * pair_n_tiles; + unsigned int _phase_x0_empty_0 = 1; + unsigned int _phase_x1_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)pair_n_tiles; + int pair_tile = tile_idx % (unsigned int)pair_n_tiles; + int off_n0 = pair_tile * 256; + int off_n1 = off_n0 + 128; + int x_row0 = batch * N + off_n0; + int x_row1 = batch * N + off_n1; + mbarrier_wait(x0_empty_addr, _phase_x0_empty_0); + _phase_x0_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x0_addr, x_tmap, 0, x_row0, 0, x0_full_addr); + mbarrier_arrive_expect_tx(x0_full_addr, 49152); + mbarrier_wait(x1_empty_addr, _phase_x1_empty_0); + _phase_x1_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x1_addr, x_tmap, 0, x_row1, 0, x1_full_addr); + mbarrier_arrive_expect_tx(x1_full_addr, 49152); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 98304); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int pair_n_tiles_1 = num_n_tiles / 2; + int num_tiles_1 = B * pair_n_tiles_1; + unsigned int _phase_x0_full_0 = 0; + unsigned int _phase_x1_full_0 = 0; + unsigned int _phase_score0_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + unsigned int _phase_score1_empty_0 = 1; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x0_full_addr, _phase_x0_full_0); + _phase_x0_full_0 ^= 1; + mbarrier_wait(x1_full_addr, _phase_x1_full_0); + _phase_x1_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score0_empty_addr, _phase_score0_empty_0); + _phase_score0_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x0_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"((tmem_score_tmem + (256))), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x0_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"((tmem_score_tmem + (256))), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x0_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"((tmem_score_tmem + (256))), "r"(1)); + elect_commit(score0_full_addr); + mbarrier_wait(score1_empty_addr, _phase_score1_empty_0); + _phase_score1_empty_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = smem_x1_addr; + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = smem_c_addr; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_3), "r"(_mma_b_lo_3), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_4 = smem_x1_addr + 16384; + int _mma_a_lo_4 = make_warp_uniform((_mma_a_addr_4 >> 4) & 0x3FFF); + int _mma_b_addr_4 = smem_c_addr + 32768; + int _mma_b_lo_4 = make_warp_uniform((_mma_b_addr_4 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_4), "r"(_mma_b_lo_4), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_5 = smem_x1_addr + 32768; + int _mma_a_lo_5 = make_warp_uniform((_mma_a_addr_5 >> 4) & 0x3FFF); + int _mma_b_addr_5 = smem_c_addr + 65536; + int _mma_b_lo_5 = make_warp_uniform((_mma_b_addr_5 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_5), "r"(_mma_b_lo_5), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score1_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x0_empty_addr); + elect_commit(x1_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int pair_n_tiles_2 = num_n_tiles / 2; + int num_tiles_2 = B * pair_n_tiles_2; + unsigned int _phase_score0_full_0 = 0; + unsigned int _phase_score1_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)pair_n_tiles_2; + int pair_tile_1 = tile_idx_2 % (unsigned int)pair_n_tiles_2; + int off_n0_1 = pair_tile_1 * 256; + int off_n1_1 = off_n0_1 + 128; + int global_n0 = off_n0_1 + (warp % 4 * 32 + lane); + int global_n1 = off_n1_1 + (warp % 4 * 32 + lane); + int out_offset0 = batch_1 * N + global_n0; + int out_offset1 = batch_1 * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best_score0 = -3.4e+38f; + int best_idx0 = 0; + float best_score1 = -3.4e+38f; + int best_idx1 = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score0_full_addr, _phase_score0_full_0); + _phase_score0_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 16) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + float csq_vals_hi[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi[0])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 4) * 4)); + float d4 = _tmem_load_0[kk + 4] - csq_vals_hi[0]; + float d5 = _tmem_load_0[kk + 5] - csq_vals_hi[1]; + float best_45 = d4; + int idx_45 = off_k_1 + score_base + kk + 4; + if (d5 > best_45) { + best_45 = d5; + idx_45 = off_k_1 + score_base + kk + 5; + } + float d6 = _tmem_load_0[kk + 6] - csq_vals_hi[2]; + float d7 = _tmem_load_0[kk + 7] - csq_vals_hi[3]; + float best_67 = d6; + int idx_67 = off_k_1 + score_base + kk + 6; + if (d7 > best_67) { + best_67 = d7; + idx_67 = off_k_1 + score_base + kk + 7; + } + float best_hi = best_45; + int idx_hi = idx_45; + if (best_67 > best_hi) { + best_hi = best_67; + idx_hi = idx_67; + } + if (best_hi > best_group) { + best_group = best_hi; + idx_group = idx_hi; + } + float csq_vals_next[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next[0])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 8) * 4)); + float d8 = _tmem_load_0[kk + 8] - csq_vals_next[0]; + float d9 = _tmem_load_0[kk + 9] - csq_vals_next[1]; + float best_89 = d8; + int idx_89 = off_k_1 + score_base + kk + 8; + if (d9 > best_89) { + best_89 = d9; + idx_89 = off_k_1 + score_base + kk + 9; + } + float d10 = _tmem_load_0[kk + 10] - csq_vals_next[2]; + float d11 = _tmem_load_0[kk + 11] - csq_vals_next[3]; + float best_1011 = d10; + int idx_1011 = off_k_1 + score_base + kk + 10; + if (d11 > best_1011) { + best_1011 = d11; + idx_1011 = off_k_1 + score_base + kk + 11; + } + float best_next = best_89; + int idx_next = idx_89; + if (best_1011 > best_next) { + best_next = best_1011; + idx_next = idx_1011; + } + float csq_vals_tail[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail[0])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk + 12) * 4)); + float d12 = _tmem_load_0[kk + 12] - csq_vals_tail[0]; + float d13 = _tmem_load_0[kk + 13] - csq_vals_tail[1]; + float best_1213 = d12; + int idx_1213 = off_k_1 + score_base + kk + 12; + if (d13 > best_1213) { + best_1213 = d13; + idx_1213 = off_k_1 + score_base + kk + 13; + } + float d14 = _tmem_load_0[kk + 14] - csq_vals_tail[2]; + float d15 = _tmem_load_0[kk + 15] - csq_vals_tail[3]; + float best_1415 = d14; + int idx_1415 = off_k_1 + score_base + kk + 14; + if (d15 > best_1415) { + best_1415 = d15; + idx_1415 = off_k_1 + score_base + kk + 15; + } + float best_tail = best_1213; + int idx_tail = idx_1213; + if (best_1415 > best_tail) { + best_tail = best_1415; + idx_tail = idx_1415; + } + if (best_tail > best_next) { + best_next = best_tail; + idx_next = idx_tail; + } + if (best_next > best_group) { + best_group = best_next; + idx_group = idx_next; + } + if (best_group > best_score0) { + best_score0 = best_group; + best_idx0 = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + 256 + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score0_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 16) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + float csq_vals_hi_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_1[0])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 4) * 4)); + float d4_1 = _tmem_load_1[kk_1 + 4] - csq_vals_hi_1[0]; + float d5_1 = _tmem_load_1[kk_1 + 5] - csq_vals_hi_1[1]; + float best_45_1 = d4_1; + int idx_45_1 = off_k_1 + score_base + kk_1 + 4; + if (d5_1 > best_45_1) { + best_45_1 = d5_1; + idx_45_1 = off_k_1 + score_base + kk_1 + 5; + } + float d6_1 = _tmem_load_1[kk_1 + 6] - csq_vals_hi_1[2]; + float d7_1 = _tmem_load_1[kk_1 + 7] - csq_vals_hi_1[3]; + float best_67_1 = d6_1; + int idx_67_1 = off_k_1 + score_base + kk_1 + 6; + if (d7_1 > best_67_1) { + best_67_1 = d7_1; + idx_67_1 = off_k_1 + score_base + kk_1 + 7; + } + float best_hi_1 = best_45_1; + int idx_hi_1 = idx_45_1; + if (best_67_1 > best_hi_1) { + best_hi_1 = best_67_1; + idx_hi_1 = idx_67_1; + } + if (best_hi_1 > best_group_1) { + best_group_1 = best_hi_1; + idx_group_1 = idx_hi_1; + } + float csq_vals_next_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_1[0])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 8) * 4)); + float d8_1 = _tmem_load_1[kk_1 + 8] - csq_vals_next_1[0]; + float d9_1 = _tmem_load_1[kk_1 + 9] - csq_vals_next_1[1]; + float best_89_1 = d8_1; + int idx_89_1 = off_k_1 + score_base + kk_1 + 8; + if (d9_1 > best_89_1) { + best_89_1 = d9_1; + idx_89_1 = off_k_1 + score_base + kk_1 + 9; + } + float d10_1 = _tmem_load_1[kk_1 + 10] - csq_vals_next_1[2]; + float d11_1 = _tmem_load_1[kk_1 + 11] - csq_vals_next_1[3]; + float best_1011_1 = d10_1; + int idx_1011_1 = off_k_1 + score_base + kk_1 + 10; + if (d11_1 > best_1011_1) { + best_1011_1 = d11_1; + idx_1011_1 = off_k_1 + score_base + kk_1 + 11; + } + float best_next_1 = best_89_1; + int idx_next_1 = idx_89_1; + if (best_1011_1 > best_next_1) { + best_next_1 = best_1011_1; + idx_next_1 = idx_1011_1; + } + float csq_vals_tail_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_1[0])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1 + 12) * 4)); + float d12_1 = _tmem_load_1[kk_1 + 12] - csq_vals_tail_1[0]; + float d13_1 = _tmem_load_1[kk_1 + 13] - csq_vals_tail_1[1]; + float best_1213_1 = d12_1; + int idx_1213_1 = off_k_1 + score_base + kk_1 + 12; + if (d13_1 > best_1213_1) { + best_1213_1 = d13_1; + idx_1213_1 = off_k_1 + score_base + kk_1 + 13; + } + float d14_1 = _tmem_load_1[kk_1 + 14] - csq_vals_tail_1[2]; + float d15_1 = _tmem_load_1[kk_1 + 15] - csq_vals_tail_1[3]; + float best_1415_1 = d14_1; + int idx_1415_1 = off_k_1 + score_base + kk_1 + 14; + if (d15_1 > best_1415_1) { + best_1415_1 = d15_1; + idx_1415_1 = off_k_1 + score_base + kk_1 + 15; + } + float best_tail_1 = best_1213_1; + int idx_tail_1 = idx_1213_1; + if (best_1415_1 > best_tail_1) { + best_tail_1 = best_1415_1; + idx_tail_1 = idx_1415_1; + } + if (best_tail_1 > best_next_1) { + best_next_1 = best_tail_1; + idx_next_1 = idx_tail_1; + } + if (best_next_1 > best_group_1) { + best_group_1 = best_next_1; + idx_group_1 = idx_next_1; + } + if (best_group_1 > best_score0) { + best_score0 = best_group_1; + best_idx0 = idx_group_1; + } + } + mbarrier_wait(score1_full_addr, _phase_score1_full_0); + _phase_score1_full_0 ^= 1; + score_base = 0; + float _tmem_load_2[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_2[0]), "=f"(_tmem_load_2[1]), "=f"(_tmem_load_2[2]), "=f"(_tmem_load_2[3]), "=f"(_tmem_load_2[4]), "=f"(_tmem_load_2[5]), "=f"(_tmem_load_2[6]), "=f"(_tmem_load_2[7]), "=f"(_tmem_load_2[8]), "=f"(_tmem_load_2[9]), "=f"(_tmem_load_2[10]), "=f"(_tmem_load_2[11]), "=f"(_tmem_load_2[12]), "=f"(_tmem_load_2[13]), "=f"(_tmem_load_2[14]), "=f"(_tmem_load_2[15]), "=f"(_tmem_load_2[16]), "=f"(_tmem_load_2[17]), "=f"(_tmem_load_2[18]), "=f"(_tmem_load_2[19]), "=f"(_tmem_load_2[20]), "=f"(_tmem_load_2[21]), "=f"(_tmem_load_2[22]), "=f"(_tmem_load_2[23]), "=f"(_tmem_load_2[24]), "=f"(_tmem_load_2[25]), "=f"(_tmem_load_2[26]), "=f"(_tmem_load_2[27]), "=f"(_tmem_load_2[28]), "=f"(_tmem_load_2[29]), "=f"(_tmem_load_2[30]), "=f"(_tmem_load_2[31]), "=f"(_tmem_load_2[32]), "=f"(_tmem_load_2[33]), "=f"(_tmem_load_2[34]), "=f"(_tmem_load_2[35]), "=f"(_tmem_load_2[36]), "=f"(_tmem_load_2[37]), "=f"(_tmem_load_2[38]), "=f"(_tmem_load_2[39]), "=f"(_tmem_load_2[40]), "=f"(_tmem_load_2[41]), "=f"(_tmem_load_2[42]), "=f"(_tmem_load_2[43]), "=f"(_tmem_load_2[44]), "=f"(_tmem_load_2[45]), "=f"(_tmem_load_2[46]), "=f"(_tmem_load_2[47]), "=f"(_tmem_load_2[48]), "=f"(_tmem_load_2[49]), "=f"(_tmem_load_2[50]), "=f"(_tmem_load_2[51]), "=f"(_tmem_load_2[52]), "=f"(_tmem_load_2[53]), "=f"(_tmem_load_2[54]), "=f"(_tmem_load_2[55]), "=f"(_tmem_load_2[56]), "=f"(_tmem_load_2[57]), "=f"(_tmem_load_2[58]), "=f"(_tmem_load_2[59]), "=f"(_tmem_load_2[60]), "=f"(_tmem_load_2[61]), "=f"(_tmem_load_2[62]), "=f"(_tmem_load_2[63]), "=f"(_tmem_load_2[64]), "=f"(_tmem_load_2[65]), "=f"(_tmem_load_2[66]), "=f"(_tmem_load_2[67]), "=f"(_tmem_load_2[68]), "=f"(_tmem_load_2[69]), "=f"(_tmem_load_2[70]), "=f"(_tmem_load_2[71]), "=f"(_tmem_load_2[72]), "=f"(_tmem_load_2[73]), "=f"(_tmem_load_2[74]), "=f"(_tmem_load_2[75]), "=f"(_tmem_load_2[76]), "=f"(_tmem_load_2[77]), "=f"(_tmem_load_2[78]), "=f"(_tmem_load_2[79]), "=f"(_tmem_load_2[80]), "=f"(_tmem_load_2[81]), "=f"(_tmem_load_2[82]), "=f"(_tmem_load_2[83]), "=f"(_tmem_load_2[84]), "=f"(_tmem_load_2[85]), "=f"(_tmem_load_2[86]), "=f"(_tmem_load_2[87]), "=f"(_tmem_load_2[88]), "=f"(_tmem_load_2[89]), "=f"(_tmem_load_2[90]), "=f"(_tmem_load_2[91]), "=f"(_tmem_load_2[92]), "=f"(_tmem_load_2[93]), "=f"(_tmem_load_2[94]), "=f"(_tmem_load_2[95]), "=f"(_tmem_load_2[96]), "=f"(_tmem_load_2[97]), "=f"(_tmem_load_2[98]), "=f"(_tmem_load_2[99]), "=f"(_tmem_load_2[100]), "=f"(_tmem_load_2[101]), "=f"(_tmem_load_2[102]), "=f"(_tmem_load_2[103]), "=f"(_tmem_load_2[104]), "=f"(_tmem_load_2[105]), "=f"(_tmem_load_2[106]), "=f"(_tmem_load_2[107]), "=f"(_tmem_load_2[108]), "=f"(_tmem_load_2[109]), "=f"(_tmem_load_2[110]), "=f"(_tmem_load_2[111]), "=f"(_tmem_load_2[112]), "=f"(_tmem_load_2[113]), "=f"(_tmem_load_2[114]), "=f"(_tmem_load_2[115]), "=f"(_tmem_load_2[116]), "=f"(_tmem_load_2[117]), "=f"(_tmem_load_2[118]), "=f"(_tmem_load_2[119]), "=f"(_tmem_load_2[120]), "=f"(_tmem_load_2[121]), "=f"(_tmem_load_2[122]), "=f"(_tmem_load_2[123]), "=f"(_tmem_load_2[124]), "=f"(_tmem_load_2[125]), "=f"(_tmem_load_2[126]), "=f"(_tmem_load_2[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk_2 = 0; kk_2 < 128; kk_2 += 16) { + float csq_vals_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_2[0])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2) * 4)); + float d0_2 = _tmem_load_2[kk_2] - csq_vals_2[0]; + float d1_2 = _tmem_load_2[kk_2 + 1] - csq_vals_2[1]; + float best_01_2 = d0_2; + int idx_01_2 = off_k_1 + score_base + kk_2; + if (d1_2 > best_01_2) { + best_01_2 = d1_2; + idx_01_2 = off_k_1 + score_base + kk_2 + 1; + } + float d2_2 = _tmem_load_2[kk_2 + 2] - csq_vals_2[2]; + float d3_2 = _tmem_load_2[kk_2 + 3] - csq_vals_2[3]; + float best_23_2 = d2_2; + int idx_23_2 = off_k_1 + score_base + kk_2 + 2; + if (d3_2 > best_23_2) { + best_23_2 = d3_2; + idx_23_2 = off_k_1 + score_base + kk_2 + 3; + } + float best_group_2 = best_01_2; + int idx_group_2 = idx_01_2; + if (best_23_2 > best_group_2) { + best_group_2 = best_23_2; + idx_group_2 = idx_23_2; + } + float csq_vals_hi_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_2[0])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 4) * 4)); + float d4_2 = _tmem_load_2[kk_2 + 4] - csq_vals_hi_2[0]; + float d5_2 = _tmem_load_2[kk_2 + 5] - csq_vals_hi_2[1]; + float best_45_2 = d4_2; + int idx_45_2 = off_k_1 + score_base + kk_2 + 4; + if (d5_2 > best_45_2) { + best_45_2 = d5_2; + idx_45_2 = off_k_1 + score_base + kk_2 + 5; + } + float d6_2 = _tmem_load_2[kk_2 + 6] - csq_vals_hi_2[2]; + float d7_2 = _tmem_load_2[kk_2 + 7] - csq_vals_hi_2[3]; + float best_67_2 = d6_2; + int idx_67_2 = off_k_1 + score_base + kk_2 + 6; + if (d7_2 > best_67_2) { + best_67_2 = d7_2; + idx_67_2 = off_k_1 + score_base + kk_2 + 7; + } + float best_hi_2 = best_45_2; + int idx_hi_2 = idx_45_2; + if (best_67_2 > best_hi_2) { + best_hi_2 = best_67_2; + idx_hi_2 = idx_67_2; + } + if (best_hi_2 > best_group_2) { + best_group_2 = best_hi_2; + idx_group_2 = idx_hi_2; + } + float csq_vals_next_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_2[0])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 8) * 4)); + float d8_2 = _tmem_load_2[kk_2 + 8] - csq_vals_next_2[0]; + float d9_2 = _tmem_load_2[kk_2 + 9] - csq_vals_next_2[1]; + float best_89_2 = d8_2; + int idx_89_2 = off_k_1 + score_base + kk_2 + 8; + if (d9_2 > best_89_2) { + best_89_2 = d9_2; + idx_89_2 = off_k_1 + score_base + kk_2 + 9; + } + float d10_2 = _tmem_load_2[kk_2 + 10] - csq_vals_next_2[2]; + float d11_2 = _tmem_load_2[kk_2 + 11] - csq_vals_next_2[3]; + float best_1011_2 = d10_2; + int idx_1011_2 = off_k_1 + score_base + kk_2 + 10; + if (d11_2 > best_1011_2) { + best_1011_2 = d11_2; + idx_1011_2 = off_k_1 + score_base + kk_2 + 11; + } + float best_next_2 = best_89_2; + int idx_next_2 = idx_89_2; + if (best_1011_2 > best_next_2) { + best_next_2 = best_1011_2; + idx_next_2 = idx_1011_2; + } + float csq_vals_tail_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_2[0])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_2[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_2 + 12) * 4)); + float d12_2 = _tmem_load_2[kk_2 + 12] - csq_vals_tail_2[0]; + float d13_2 = _tmem_load_2[kk_2 + 13] - csq_vals_tail_2[1]; + float best_1213_2 = d12_2; + int idx_1213_2 = off_k_1 + score_base + kk_2 + 12; + if (d13_2 > best_1213_2) { + best_1213_2 = d13_2; + idx_1213_2 = off_k_1 + score_base + kk_2 + 13; + } + float d14_2 = _tmem_load_2[kk_2 + 14] - csq_vals_tail_2[2]; + float d15_2 = _tmem_load_2[kk_2 + 15] - csq_vals_tail_2[3]; + float best_1415_2 = d14_2; + int idx_1415_2 = off_k_1 + score_base + kk_2 + 14; + if (d15_2 > best_1415_2) { + best_1415_2 = d15_2; + idx_1415_2 = off_k_1 + score_base + kk_2 + 15; + } + float best_tail_2 = best_1213_2; + int idx_tail_2 = idx_1213_2; + if (best_1415_2 > best_tail_2) { + best_tail_2 = best_1415_2; + idx_tail_2 = idx_1415_2; + } + if (best_tail_2 > best_next_2) { + best_next_2 = best_tail_2; + idx_next_2 = idx_tail_2; + } + if (best_next_2 > best_group_2) { + best_group_2 = best_next_2; + idx_group_2 = idx_next_2; + } + if (best_group_2 > best_score1) { + best_score1 = best_group_2; + best_idx1 = idx_group_2; + } + } + score_base = 128; + float _tmem_load_3[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_3[0]), "=f"(_tmem_load_3[1]), "=f"(_tmem_load_3[2]), "=f"(_tmem_load_3[3]), "=f"(_tmem_load_3[4]), "=f"(_tmem_load_3[5]), "=f"(_tmem_load_3[6]), "=f"(_tmem_load_3[7]), "=f"(_tmem_load_3[8]), "=f"(_tmem_load_3[9]), "=f"(_tmem_load_3[10]), "=f"(_tmem_load_3[11]), "=f"(_tmem_load_3[12]), "=f"(_tmem_load_3[13]), "=f"(_tmem_load_3[14]), "=f"(_tmem_load_3[15]), "=f"(_tmem_load_3[16]), "=f"(_tmem_load_3[17]), "=f"(_tmem_load_3[18]), "=f"(_tmem_load_3[19]), "=f"(_tmem_load_3[20]), "=f"(_tmem_load_3[21]), "=f"(_tmem_load_3[22]), "=f"(_tmem_load_3[23]), "=f"(_tmem_load_3[24]), "=f"(_tmem_load_3[25]), "=f"(_tmem_load_3[26]), "=f"(_tmem_load_3[27]), "=f"(_tmem_load_3[28]), "=f"(_tmem_load_3[29]), "=f"(_tmem_load_3[30]), "=f"(_tmem_load_3[31]), "=f"(_tmem_load_3[32]), "=f"(_tmem_load_3[33]), "=f"(_tmem_load_3[34]), "=f"(_tmem_load_3[35]), "=f"(_tmem_load_3[36]), "=f"(_tmem_load_3[37]), "=f"(_tmem_load_3[38]), "=f"(_tmem_load_3[39]), "=f"(_tmem_load_3[40]), "=f"(_tmem_load_3[41]), "=f"(_tmem_load_3[42]), "=f"(_tmem_load_3[43]), "=f"(_tmem_load_3[44]), "=f"(_tmem_load_3[45]), "=f"(_tmem_load_3[46]), "=f"(_tmem_load_3[47]), "=f"(_tmem_load_3[48]), "=f"(_tmem_load_3[49]), "=f"(_tmem_load_3[50]), "=f"(_tmem_load_3[51]), "=f"(_tmem_load_3[52]), "=f"(_tmem_load_3[53]), "=f"(_tmem_load_3[54]), "=f"(_tmem_load_3[55]), "=f"(_tmem_load_3[56]), "=f"(_tmem_load_3[57]), "=f"(_tmem_load_3[58]), "=f"(_tmem_load_3[59]), "=f"(_tmem_load_3[60]), "=f"(_tmem_load_3[61]), "=f"(_tmem_load_3[62]), "=f"(_tmem_load_3[63]), "=f"(_tmem_load_3[64]), "=f"(_tmem_load_3[65]), "=f"(_tmem_load_3[66]), "=f"(_tmem_load_3[67]), "=f"(_tmem_load_3[68]), "=f"(_tmem_load_3[69]), "=f"(_tmem_load_3[70]), "=f"(_tmem_load_3[71]), "=f"(_tmem_load_3[72]), "=f"(_tmem_load_3[73]), "=f"(_tmem_load_3[74]), "=f"(_tmem_load_3[75]), "=f"(_tmem_load_3[76]), "=f"(_tmem_load_3[77]), "=f"(_tmem_load_3[78]), "=f"(_tmem_load_3[79]), "=f"(_tmem_load_3[80]), "=f"(_tmem_load_3[81]), "=f"(_tmem_load_3[82]), "=f"(_tmem_load_3[83]), "=f"(_tmem_load_3[84]), "=f"(_tmem_load_3[85]), "=f"(_tmem_load_3[86]), "=f"(_tmem_load_3[87]), "=f"(_tmem_load_3[88]), "=f"(_tmem_load_3[89]), "=f"(_tmem_load_3[90]), "=f"(_tmem_load_3[91]), "=f"(_tmem_load_3[92]), "=f"(_tmem_load_3[93]), "=f"(_tmem_load_3[94]), "=f"(_tmem_load_3[95]), "=f"(_tmem_load_3[96]), "=f"(_tmem_load_3[97]), "=f"(_tmem_load_3[98]), "=f"(_tmem_load_3[99]), "=f"(_tmem_load_3[100]), "=f"(_tmem_load_3[101]), "=f"(_tmem_load_3[102]), "=f"(_tmem_load_3[103]), "=f"(_tmem_load_3[104]), "=f"(_tmem_load_3[105]), "=f"(_tmem_load_3[106]), "=f"(_tmem_load_3[107]), "=f"(_tmem_load_3[108]), "=f"(_tmem_load_3[109]), "=f"(_tmem_load_3[110]), "=f"(_tmem_load_3[111]), "=f"(_tmem_load_3[112]), "=f"(_tmem_load_3[113]), "=f"(_tmem_load_3[114]), "=f"(_tmem_load_3[115]), "=f"(_tmem_load_3[116]), "=f"(_tmem_load_3[117]), "=f"(_tmem_load_3[118]), "=f"(_tmem_load_3[119]), "=f"(_tmem_load_3[120]), "=f"(_tmem_load_3[121]), "=f"(_tmem_load_3[122]), "=f"(_tmem_load_3[123]), "=f"(_tmem_load_3[124]), "=f"(_tmem_load_3[125]), "=f"(_tmem_load_3[126]), "=f"(_tmem_load_3[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score1_empty_addr); + } + #pragma unroll + for (int kk_3 = 0; kk_3 < 128; kk_3 += 16) { + float csq_vals_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_3[0])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3) * 4)); + float d0_3 = _tmem_load_3[kk_3] - csq_vals_3[0]; + float d1_3 = _tmem_load_3[kk_3 + 1] - csq_vals_3[1]; + float best_01_3 = d0_3; + int idx_01_3 = off_k_1 + score_base + kk_3; + if (d1_3 > best_01_3) { + best_01_3 = d1_3; + idx_01_3 = off_k_1 + score_base + kk_3 + 1; + } + float d2_3 = _tmem_load_3[kk_3 + 2] - csq_vals_3[2]; + float d3_3 = _tmem_load_3[kk_3 + 3] - csq_vals_3[3]; + float best_23_3 = d2_3; + int idx_23_3 = off_k_1 + score_base + kk_3 + 2; + if (d3_3 > best_23_3) { + best_23_3 = d3_3; + idx_23_3 = off_k_1 + score_base + kk_3 + 3; + } + float best_group_3 = best_01_3; + int idx_group_3 = idx_01_3; + if (best_23_3 > best_group_3) { + best_group_3 = best_23_3; + idx_group_3 = idx_23_3; + } + float csq_vals_hi_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_hi_3[0])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_hi_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 4) * 4)); + float d4_3 = _tmem_load_3[kk_3 + 4] - csq_vals_hi_3[0]; + float d5_3 = _tmem_load_3[kk_3 + 5] - csq_vals_hi_3[1]; + float best_45_3 = d4_3; + int idx_45_3 = off_k_1 + score_base + kk_3 + 4; + if (d5_3 > best_45_3) { + best_45_3 = d5_3; + idx_45_3 = off_k_1 + score_base + kk_3 + 5; + } + float d6_3 = _tmem_load_3[kk_3 + 6] - csq_vals_hi_3[2]; + float d7_3 = _tmem_load_3[kk_3 + 7] - csq_vals_hi_3[3]; + float best_67_3 = d6_3; + int idx_67_3 = off_k_1 + score_base + kk_3 + 6; + if (d7_3 > best_67_3) { + best_67_3 = d7_3; + idx_67_3 = off_k_1 + score_base + kk_3 + 7; + } + float best_hi_3 = best_45_3; + int idx_hi_3 = idx_45_3; + if (best_67_3 > best_hi_3) { + best_hi_3 = best_67_3; + idx_hi_3 = idx_67_3; + } + if (best_hi_3 > best_group_3) { + best_group_3 = best_hi_3; + idx_group_3 = idx_hi_3; + } + float csq_vals_next_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_next_3[0])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_next_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 8) * 4)); + float d8_3 = _tmem_load_3[kk_3 + 8] - csq_vals_next_3[0]; + float d9_3 = _tmem_load_3[kk_3 + 9] - csq_vals_next_3[1]; + float best_89_3 = d8_3; + int idx_89_3 = off_k_1 + score_base + kk_3 + 8; + if (d9_3 > best_89_3) { + best_89_3 = d9_3; + idx_89_3 = off_k_1 + score_base + kk_3 + 9; + } + float d10_3 = _tmem_load_3[kk_3 + 10] - csq_vals_next_3[2]; + float d11_3 = _tmem_load_3[kk_3 + 11] - csq_vals_next_3[3]; + float best_1011_3 = d10_3; + int idx_1011_3 = off_k_1 + score_base + kk_3 + 10; + if (d11_3 > best_1011_3) { + best_1011_3 = d11_3; + idx_1011_3 = off_k_1 + score_base + kk_3 + 11; + } + float best_next_3 = best_89_3; + int idx_next_3 = idx_89_3; + if (best_1011_3 > best_next_3) { + best_next_3 = best_1011_3; + idx_next_3 = idx_1011_3; + } + float csq_vals_tail_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_tail_3[0])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_tail_3[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_3 + 12) * 4)); + float d12_3 = _tmem_load_3[kk_3 + 12] - csq_vals_tail_3[0]; + float d13_3 = _tmem_load_3[kk_3 + 13] - csq_vals_tail_3[1]; + float best_1213_3 = d12_3; + int idx_1213_3 = off_k_1 + score_base + kk_3 + 12; + if (d13_3 > best_1213_3) { + best_1213_3 = d13_3; + idx_1213_3 = off_k_1 + score_base + kk_3 + 13; + } + float d14_3 = _tmem_load_3[kk_3 + 14] - csq_vals_tail_3[2]; + float d15_3 = _tmem_load_3[kk_3 + 15] - csq_vals_tail_3[3]; + float best_1415_3 = d14_3; + int idx_1415_3 = off_k_1 + score_base + kk_3 + 14; + if (d15_3 > best_1415_3) { + best_1415_3 = d15_3; + idx_1415_3 = off_k_1 + score_base + kk_3 + 15; + } + float best_tail_3 = best_1213_3; + int idx_tail_3 = idx_1213_3; + if (best_1415_3 > best_tail_3) { + best_tail_3 = best_1415_3; + idx_tail_3 = idx_1415_3; + } + if (best_tail_3 > best_next_3) { + best_next_3 = best_tail_3; + idx_next_3 = idx_tail_3; + } + if (best_next_3 > best_group_3) { + best_group_3 = best_next_3; + idx_group_3 = idx_next_3; + } + if (best_group_3 > best_score1) { + best_score1 = best_group_3; + best_idx1 = idx_group_3; + } + } + } + if (global_n0 < N) { + *((int*)(out + out_offset0)) = best_idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = best_idx1; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0344.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0344.cu new file mode 100644 index 00000000..7070fa90 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0344.cu @@ -0,0 +1,600 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 49152 +#define SMEM_SMEM_X_STRIDE 49152 +#define SMEM_SMEM_C_OFF 50176 +#define SMEM_SMEM_C_STAGE_BYTES 98304 +#define SMEM_SMEM_C_STRIDE 98304 +#define SMEM_SMEM_CSQ_OFF 148480 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 149504 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 50176); + const int smem_c_addr = smem + 50176; + float* smem_csq = reinterpret_cast(smem_raw + 148480); + const int smem_csq_addr = smem + 148480; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, 0, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 49152); + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, 0, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 98304); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(0)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = smem_x_addr + 16384; + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = smem_c_addr + 32768; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_1), "r"(_mma_b_lo_1), "r"(tmem_score_tmem), "r"(1)); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = smem_x_addr + 32768; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = smem_c_addr + 65536; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_2), "r"(_mma_b_lo_2), "r"(tmem_score_tmem), "r"(1)); + elect_commit(score_full_addr); + elect_commit(c_empty_addr); + } + elect_commit(x_empty_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0345.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0345.cu new file mode 100644 index 00000000..3f16db31 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0345.cu @@ -0,0 +1,525 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 4096 +#define SMEM_SMEM_X_STRIDE 4096 +#define SMEM_SMEM_C_OFF 5120 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 21504 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 22528 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d480_splitk_partial_d32k256_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, float* __restrict__ partial_scores, int* __restrict__ partial_indices, int B, int N, int D, int K, int num_n_tiles, int K_tiles, int K_slices) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 5120); + const int smem_c_addr = smem + 5120; + float* smem_csq = reinterpret_cast(smem_raw + 21504); + const int smem_csq_addr = smem + 21504; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: compute ---- + if (warp <= 3) { + { // compute_main + int total_work = B * num_n_tiles * K_slices; + int compute_warp = warp; + int lane_pair = lane % 4; + int row_origin = compute_warp * 16; + int row_lane_base = row_origin + lane / 4; + int row0 = row_lane_base; + int row1 = row_lane_base + 8; + int compute_tid = compute_warp * 32 + lane; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_slice = work_idx % (unsigned int)K_slices; + int point_tile_idx = work_idx / (unsigned int)K_slices; + int batch = point_tile_idx / num_n_tiles; + int slice_k_start = iter_slice; + int csq_smem_addr = smem_csq_addr; + float best0 = -3.4e+38f; + float best1 = -3.4e+38f; + int idx0 = slice_k_start * 256; + int idx1 = idx0; + #pragma unroll 1 + for (int local_k = 0; local_k < 1; local_k++) { + int iter_k = slice_k_start + local_k; + int off_k = iter_k * 256; + if ((float)compute_tid < 64.0f) { + int csq_base = compute_tid * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch * K + off_k + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + #pragma unroll + for (int score_base = 0; score_base < 256; score_base += 128) { + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, [%64];" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])), "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) + : "r"(taddr + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int rep = 0; rep < 16.0f; rep++) { + int local_reg = rep * 4; + int col_base = score_base + rep * 8 + lane_pair * 2; + float csq0 = smem_csq[col_base]; + float csq1 = smem_csq[col_base + 1]; + float d0 = _tmem_load_0[local_reg] - csq0; + if (d0 > best0) { + best0 = d0; + idx0 = off_k + col_base; + } + float d1 = _tmem_load_0[local_reg + 1] - csq1; + if (d1 > best0) { + best0 = d1; + idx0 = off_k + col_base + 1; + } + float d2 = _tmem_load_0[local_reg + 2] - csq0; + if (d2 > best1) { + best1 = d2; + idx1 = off_k + col_base; + } + float d3 = _tmem_load_0[local_reg + 3] - csq1; + if (d3 > best1) { + best1 = d3; + idx1 = off_k + col_base + 1; + } + } + } + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best0, 1); + float peer0 = _shfl_xor_0; + int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, idx0, 1); + int peer0_idx = _shfl_xor_1; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer0_idx; + } + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, best0, 2); + peer0 = _shfl_xor_2; + int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, idx0, 2); + peer0_idx = _shfl_xor_3; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer0_idx; + } + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, best1, 1); + float peer1 = _shfl_xor_4; + int _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, idx1, 1); + int peer1_idx = _shfl_xor_5; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer1_idx; + } + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, best1, 2); + peer1 = _shfl_xor_6; + int _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, idx1, 2); + peer1_idx = _shfl_xor_7; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer1_idx; + } + if (lane_pair == 0) { + int partial_offset0 = work_idx * 64 + (unsigned int)row0; + *((float*)(partial_scores + partial_offset0)) = best0; + *((int*)(partial_indices + partial_offset0)) = idx0; + int partial_offset1 = work_idx * 64 + (unsigned int)row1; + *((float*)(partial_scores + partial_offset1)) = best1; + *((int*)(partial_indices + partial_offset1)) = idx1; + } + } + } + // ---- Role: load ---- + } else if (warp == 4) { + { // load_main + int total_work_1 = B * num_n_tiles * K_slices; + int feature_tiles = D / 32; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_1; work_idx_1 += num_bids) { + int iter_slice_1 = work_idx_1 % (unsigned int)K_slices; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_slices; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int n_tile = point_tile_idx_1 % num_n_tiles; + int off_n = n_tile * 64; + int x_row = batch_1 * N + off_n; + int slice_k_start_1 = iter_slice_1; + #pragma unroll 1 + for (int local_k_1 = 0; local_k_1 < 1; local_k_1++) { + int iter_k_1 = slice_k_start_1 + local_k_1; + int off_k_1 = iter_k_1 * 256; + int c_row = batch_1 * K + off_k_1; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 4096); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 5) { + { // mma_main + int total_work_2 = B * num_n_tiles * K_slices; + int feature_tiles_1 = D / 32; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_2 = bid; work_idx_2 < total_work_2; work_idx_2 += num_bids) { + int iter_slice_2 = work_idx_2 % (unsigned int)K_slices; + int slice_k_start_2 = iter_slice_2; + #pragma unroll 1 + for (int local_k_2 = 0; local_k_2 < 1; local_k_2++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 71304336;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0346.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0346.cu new file mode 100644 index 00000000..c242dff7 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0346.cu @@ -0,0 +1,81 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d480_splitk_reduce_d32k256_v1(float* __restrict__ partial_scores, int* __restrict__ partial_indices, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_slices) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid / 4; + int row_lane = tid % 4; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 64 + row; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_slice = row_lane; iter_slice < K_slices; iter_slice += 4) { + int partial_offset = (point_tile_idx * (unsigned int)K_slices + (unsigned int)iter_slice) * 64 + (unsigned int)row; + float score = partial_scores[partial_offset]; + int idx = partial_indices[partial_offset]; + if (score > best_score) { + best_score = score; + best_idx = idx; + } + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best_score, 1); + float peer_score = _shfl_xor_0; + int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, best_idx, 1); + int peer_idx = _shfl_xor_1; + if (peer_score > best_score) { + best_score = peer_score; + best_idx = peer_idx; + } + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, best_score, 2); + peer_score = _shfl_xor_2; + int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, best_idx, 2); + peer_idx = _shfl_xor_3; + if (peer_score > best_score) { + best_score = peer_score; + best_idx = peer_idx; + } + if (row_lane == 0) { + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = best_idx; + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0347.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0347.cu new file mode 100644 index 00000000..190f65dc --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0347.cu @@ -0,0 +1,520 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 25600 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 26624 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d416_exactd_splitd_a4a579d1_v2(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 25600); + const int smem_csq_addr = smem + 25600; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int iter_d = 0; iter_d < D / 32; iter_d++) { + int d_group = iter_d; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, d_group, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, d_group, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int iter_d_1 = 0; iter_d_1 < D / 32; iter_d_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((iter_d_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0348.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0348.cu new file mode 100644 index 00000000..edfe179b --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0348.cu @@ -0,0 +1,512 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 25600 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 26624 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, float* __restrict__ partial_scores, int* __restrict__ partial_indices, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 25600); + const int smem_csq_addr = smem + 25600; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int total_work = B * num_n_tiles * K_tiles; + int feature_tiles = 9; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int work_idx = bid; work_idx < total_work; work_idx += num_bids) { + int iter_k = work_idx % (unsigned int)K_tiles; + int point_tile_idx = work_idx / (unsigned int)K_tiles; + int batch = point_tile_idx / num_n_tiles; + int n_tile = point_tile_idx % num_n_tiles; + int off_n = n_tile * 128; + int off_k = iter_k * 256; + int x_row = batch * N + off_n; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int feat_tile = 0; feat_tile < feature_tiles; feat_tile++) { + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, feat_tile, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, feat_tile, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int total_work_1 = B * num_n_tiles * K_tiles; + int feature_tiles_1 = 9; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int _work_idx = bid; _work_idx < total_work_1; _work_idx += num_bids) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int feat_tile_1 = 0; feat_tile_1 < feature_tiles_1; feat_tile_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((feat_tile_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int total_work_2 = B * num_n_tiles * K_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int work_idx_1 = bid; work_idx_1 < total_work_2; work_idx_1 += num_bids) { + int iter_k_1 = work_idx_1 % (unsigned int)K_tiles; + int point_tile_idx_1 = work_idx_1 / (unsigned int)K_tiles; + int batch_1 = point_tile_idx_1 / num_n_tiles; + int off_k_1 = iter_k_1 * 256; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = off_k_1; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + int partial_offset = work_idx_1 * 128 + (unsigned int)(warp % 4 * 32 + lane); + *((float*)(partial_scores + partial_offset)) = best_score; + *((int*)(partial_indices + partial_offset)) = best_idx; + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0349.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0349.cu new file mode 100644 index 00000000..0249cc88 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0349.cu @@ -0,0 +1,62 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 128 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce(float* __restrict__ partial_scores, int* __restrict__ partial_indices, int* __restrict__ out, int B, int N, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int row = tid; + int total_point_tiles = B * num_n_tiles; + #pragma unroll 1 + for (unsigned int point_tile_idx = bid; point_tile_idx < total_point_tiles; point_tile_idx += num_bids) { + int batch = point_tile_idx / (unsigned int)num_n_tiles; + int n_tile = point_tile_idx % (unsigned int)num_n_tiles; + int global_n = n_tile * 128 + row; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int partial_offset = (point_tile_idx * (unsigned int)K_tiles + (unsigned int)iter_k) * 128 + (unsigned int)row; + float score = partial_scores[partial_offset]; + int idx = partial_indices[partial_offset]; + if (score > best_score) { + best_score = score; + best_idx = idx; + } + } + int out_offset = batch * N + global_n; + *((int*)(out + out_offset)) = best_idx; + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0350.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0350.cu new file mode 100644 index 00000000..f2293d4d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0350.cu @@ -0,0 +1,520 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 8192 +#define SMEM_SMEM_X_STRIDE 8192 +#define SMEM_SMEM_C_OFF 9216 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 25600 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 26624 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d288_exactd_a532_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_c_addr = smem + 9216; + float* smem_csq = reinterpret_cast(smem_raw + 25600); + const int smem_csq_addr = smem + 25600; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int num_tiles = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 128; + int x_row = batch * N + off_n; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + int c_row = batch * K + off_k; + #pragma unroll 1 + for (int iter_d = 0; iter_d < 9; iter_d++) { + int d_group = iter_d; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, d_group, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 8192); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, d_group, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 1) { + { // mma_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int iter_d_1 = 0; iter_d_1 < 9; iter_d_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((iter_d_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 138413200;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 2 && warp <= 5) { + { // compute_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + int batch_1 = tile_idx_2 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_2 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 128; + int global_n = off_n_1 + (warp % 4 * 32 + lane); + int out_offset = batch_1 * N + global_n; + int csq_smem_addr = smem_csq_addr; + float best_score = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + int off_k_1 = iter_k_2 * 256; + if (warp % 4 * 32 + lane < 64) { + int csq_base = (warp % 4 * 32 + lane) * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch_1 * K + off_k_1 + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + int score_base = 0; + float _tmem_load_0[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]), "=f"(_tmem_load_0[64]), "=f"(_tmem_load_0[65]), "=f"(_tmem_load_0[66]), "=f"(_tmem_load_0[67]), "=f"(_tmem_load_0[68]), "=f"(_tmem_load_0[69]), "=f"(_tmem_load_0[70]), "=f"(_tmem_load_0[71]), "=f"(_tmem_load_0[72]), "=f"(_tmem_load_0[73]), "=f"(_tmem_load_0[74]), "=f"(_tmem_load_0[75]), "=f"(_tmem_load_0[76]), "=f"(_tmem_load_0[77]), "=f"(_tmem_load_0[78]), "=f"(_tmem_load_0[79]), "=f"(_tmem_load_0[80]), "=f"(_tmem_load_0[81]), "=f"(_tmem_load_0[82]), "=f"(_tmem_load_0[83]), "=f"(_tmem_load_0[84]), "=f"(_tmem_load_0[85]), "=f"(_tmem_load_0[86]), "=f"(_tmem_load_0[87]), "=f"(_tmem_load_0[88]), "=f"(_tmem_load_0[89]), "=f"(_tmem_load_0[90]), "=f"(_tmem_load_0[91]), "=f"(_tmem_load_0[92]), "=f"(_tmem_load_0[93]), "=f"(_tmem_load_0[94]), "=f"(_tmem_load_0[95]), "=f"(_tmem_load_0[96]), "=f"(_tmem_load_0[97]), "=f"(_tmem_load_0[98]), "=f"(_tmem_load_0[99]), "=f"(_tmem_load_0[100]), "=f"(_tmem_load_0[101]), "=f"(_tmem_load_0[102]), "=f"(_tmem_load_0[103]), "=f"(_tmem_load_0[104]), "=f"(_tmem_load_0[105]), "=f"(_tmem_load_0[106]), "=f"(_tmem_load_0[107]), "=f"(_tmem_load_0[108]), "=f"(_tmem_load_0[109]), "=f"(_tmem_load_0[110]), "=f"(_tmem_load_0[111]), "=f"(_tmem_load_0[112]), "=f"(_tmem_load_0[113]), "=f"(_tmem_load_0[114]), "=f"(_tmem_load_0[115]), "=f"(_tmem_load_0[116]), "=f"(_tmem_load_0[117]), "=f"(_tmem_load_0[118]), "=f"(_tmem_load_0[119]), "=f"(_tmem_load_0[120]), "=f"(_tmem_load_0[121]), "=f"(_tmem_load_0[122]), "=f"(_tmem_load_0[123]), "=f"(_tmem_load_0[124]), "=f"(_tmem_load_0[125]), "=f"(_tmem_load_0[126]), "=f"(_tmem_load_0[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int kk = 0; kk < 128; kk += 4) { + float csq_vals[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals[0])), "=r"(*reinterpret_cast(&csq_vals[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk) * 4)); + float d0 = _tmem_load_0[kk] - csq_vals[0]; + float d1 = _tmem_load_0[kk + 1] - csq_vals[1]; + float best_01 = d0; + int idx_01 = off_k_1 + score_base + kk; + if (d1 > best_01) { + best_01 = d1; + idx_01 = off_k_1 + score_base + kk + 1; + } + float d2 = _tmem_load_0[kk + 2] - csq_vals[2]; + float d3 = _tmem_load_0[kk + 3] - csq_vals[3]; + float best_23 = d2; + int idx_23 = off_k_1 + score_base + kk + 2; + if (d3 > best_23) { + best_23 = d3; + idx_23 = off_k_1 + score_base + kk + 3; + } + float best_group = best_01; + int idx_group = idx_01; + if (best_23 > best_group) { + best_group = best_23; + idx_group = idx_23; + } + if (best_group > best_score) { + best_score = best_group; + best_idx = idx_group; + } + } + score_base = 128; + float _tmem_load_1[128]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x128.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63, %64, %65, %66, %67, %68, %69, %70, %71, %72, %73, %74, %75, %76, %77, %78, %79, %80, %81, %82, %83, %84, %85, %86, %87, %88, %89, %90, %91, %92, %93, %94, %95, %96, %97, %98, %99, %100, %101, %102, %103, %104, %105, %106, %107, %108, %109, %110, %111, %112, %113, %114, %115, %116, %117, %118, %119, %120, %121, %122, %123, %124, %125, %126, %127}, [%128];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]), "=f"(_tmem_load_1[64]), "=f"(_tmem_load_1[65]), "=f"(_tmem_load_1[66]), "=f"(_tmem_load_1[67]), "=f"(_tmem_load_1[68]), "=f"(_tmem_load_1[69]), "=f"(_tmem_load_1[70]), "=f"(_tmem_load_1[71]), "=f"(_tmem_load_1[72]), "=f"(_tmem_load_1[73]), "=f"(_tmem_load_1[74]), "=f"(_tmem_load_1[75]), "=f"(_tmem_load_1[76]), "=f"(_tmem_load_1[77]), "=f"(_tmem_load_1[78]), "=f"(_tmem_load_1[79]), "=f"(_tmem_load_1[80]), "=f"(_tmem_load_1[81]), "=f"(_tmem_load_1[82]), "=f"(_tmem_load_1[83]), "=f"(_tmem_load_1[84]), "=f"(_tmem_load_1[85]), "=f"(_tmem_load_1[86]), "=f"(_tmem_load_1[87]), "=f"(_tmem_load_1[88]), "=f"(_tmem_load_1[89]), "=f"(_tmem_load_1[90]), "=f"(_tmem_load_1[91]), "=f"(_tmem_load_1[92]), "=f"(_tmem_load_1[93]), "=f"(_tmem_load_1[94]), "=f"(_tmem_load_1[95]), "=f"(_tmem_load_1[96]), "=f"(_tmem_load_1[97]), "=f"(_tmem_load_1[98]), "=f"(_tmem_load_1[99]), "=f"(_tmem_load_1[100]), "=f"(_tmem_load_1[101]), "=f"(_tmem_load_1[102]), "=f"(_tmem_load_1[103]), "=f"(_tmem_load_1[104]), "=f"(_tmem_load_1[105]), "=f"(_tmem_load_1[106]), "=f"(_tmem_load_1[107]), "=f"(_tmem_load_1[108]), "=f"(_tmem_load_1[109]), "=f"(_tmem_load_1[110]), "=f"(_tmem_load_1[111]), "=f"(_tmem_load_1[112]), "=f"(_tmem_load_1[113]), "=f"(_tmem_load_1[114]), "=f"(_tmem_load_1[115]), "=f"(_tmem_load_1[116]), "=f"(_tmem_load_1[117]), "=f"(_tmem_load_1[118]), "=f"(_tmem_load_1[119]), "=f"(_tmem_load_1[120]), "=f"(_tmem_load_1[121]), "=f"(_tmem_load_1[122]), "=f"(_tmem_load_1[123]), "=f"(_tmem_load_1[124]), "=f"(_tmem_load_1[125]), "=f"(_tmem_load_1[126]), "=f"(_tmem_load_1[127]) + : "r"(taddr + (unsigned int)(warp % 4 * 32 << 16) + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 128; kk_1 += 4) { + float csq_vals_1[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&csq_vals_1[0])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 1])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 2])), "=r"(*reinterpret_cast(&csq_vals_1[(0) + 3])) + : "r"(csq_smem_addr + (score_base + kk_1) * 4)); + float d0_1 = _tmem_load_1[kk_1] - csq_vals_1[0]; + float d1_1 = _tmem_load_1[kk_1 + 1] - csq_vals_1[1]; + float best_01_1 = d0_1; + int idx_01_1 = off_k_1 + score_base + kk_1; + if (d1_1 > best_01_1) { + best_01_1 = d1_1; + idx_01_1 = off_k_1 + score_base + kk_1 + 1; + } + float d2_1 = _tmem_load_1[kk_1 + 2] - csq_vals_1[2]; + float d3_1 = _tmem_load_1[kk_1 + 3] - csq_vals_1[3]; + float best_23_1 = d2_1; + int idx_23_1 = off_k_1 + score_base + kk_1 + 2; + if (d3_1 > best_23_1) { + best_23_1 = d3_1; + idx_23_1 = off_k_1 + score_base + kk_1 + 3; + } + float best_group_1 = best_01_1; + int idx_group_1 = idx_01_1; + if (best_23_1 > best_group_1) { + best_group_1 = best_23_1; + idx_group_1 = idx_23_1; + } + if (best_group_1 > best_score) { + best_score = best_group_1; + best_idx = idx_group_1; + } + } + } + if (global_n < N) { + *((int*)(out + out_offset)) = best_idx; + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0351.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0351.cu new file mode 100644 index 00000000..fcc7623c --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0351.cu @@ -0,0 +1,520 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_SCORE_TMEM_OFFSET 0 +#define NUM_MAIN_STAGES 1 +#define SMEM_SMEM_X_OFF 1024 +#define SMEM_SMEM_X_STAGE_BYTES 4096 +#define SMEM_SMEM_X_STRIDE 4096 +#define SMEM_SMEM_C_OFF 5120 +#define SMEM_SMEM_C_STAGE_BYTES 16384 +#define SMEM_SMEM_C_STRIDE 16384 +#define SMEM_SMEM_CSQ_OFF 21504 +#define SMEM_SMEM_CSQ_STAGE_BYTES 1024 +#define SMEM_SMEM_CSQ_STRIDE 1024 +#define SMEM_TOTAL 22528 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(192) void +kernel_flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ x_sq, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_x = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_x_addr = smem + 1024; + __nv_bfloat16* smem_c = reinterpret_cast<__nv_bfloat16*>(smem_raw + 5120); + const int smem_c_addr = smem + 5120; + float* smem_csq = reinterpret_cast(smem_raw + 21504); + const int smem_csq_addr = smem + 21504; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // x_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // score_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // score_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 40, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 48); + if (warp == 0) { + int _tmem_hold = smem + 48; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define x_empty_addr (mbar_base + 8) + #define c_full_addr (mbar_base + 16) + #define c_empty_addr (mbar_base + 24) + #define score_full_addr (mbar_base + 32) + #define score_empty_addr (mbar_base + 40) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_score_tmem = taddr; + + // ---- Role: compute ---- + if (warp <= 3) { + { // compute_main + int num_tiles = B * num_n_tiles; + int compute_warp = warp; + int lane_pair = lane % 4; + int row_origin = compute_warp * 16; + int row_lane_base = row_origin + lane / 4; + int row0 = row_lane_base; + int row1 = row_lane_base + 8; + int compute_tid = compute_warp * 32 + lane; + unsigned int _phase_score_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx = bid; tile_idx < num_tiles; tile_idx += num_bids) { + int batch = tile_idx / (unsigned int)num_n_tiles; + int n_tile = tile_idx % (unsigned int)num_n_tiles; + int off_n = n_tile * 64; + int global_n0 = off_n + row0; + int global_n1 = off_n + row1; + int out_offset0 = batch * N + global_n0; + int out_offset1 = batch * N + global_n1; + int csq_smem_addr = smem_csq_addr; + float best0 = -3.4e+38f; + float best1 = -3.4e+38f; + int idx0 = 0; + int idx1 = 0; + #pragma unroll 1 + for (int iter_k = 0; iter_k < K_tiles; iter_k++) { + int off_k = iter_k * 256; + if ((float)compute_tid < 64.0f) { + int csq_base = compute_tid * 4; + float _vec_load_0[4]; + { + float4 _v4 = *reinterpret_cast(c_sq + batch * K + off_k + csq_base); + _vec_load_0[0 + 0] = _v4.x; + _vec_load_0[0 + 1] = _v4.y; + _vec_load_0[0 + 2] = _v4.z; + _vec_load_0[0 + 3] = _v4.w; + } + float csq_half[4]; + csq_half[0] = 0.5f * _vec_load_0[0]; + csq_half[1] = 0.5f * _vec_load_0[1]; + csq_half[2] = 0.5f * _vec_load_0[2]; + csq_half[3] = 0.5f * _vec_load_0[3]; + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(csq_smem_addr + csq_base * 4), "f"(csq_half[0]), "f"(csq_half[1]), "f"(csq_half[2]), "f"(csq_half[3]) : "memory"); + } + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + mbarrier_wait(score_full_addr, _phase_score_full_0); + _phase_score_full_0 ^= 1; + #pragma unroll + for (int score_base = 0; score_base < 256; score_base += 128) { + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, [%64];" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])), "=r"(*reinterpret_cast(&_tmem_load_0[32])), "=r"(*reinterpret_cast(&_tmem_load_0[33])), "=r"(*reinterpret_cast(&_tmem_load_0[34])), "=r"(*reinterpret_cast(&_tmem_load_0[35])), "=r"(*reinterpret_cast(&_tmem_load_0[36])), "=r"(*reinterpret_cast(&_tmem_load_0[37])), "=r"(*reinterpret_cast(&_tmem_load_0[38])), "=r"(*reinterpret_cast(&_tmem_load_0[39])), "=r"(*reinterpret_cast(&_tmem_load_0[40])), "=r"(*reinterpret_cast(&_tmem_load_0[41])), "=r"(*reinterpret_cast(&_tmem_load_0[42])), "=r"(*reinterpret_cast(&_tmem_load_0[43])), "=r"(*reinterpret_cast(&_tmem_load_0[44])), "=r"(*reinterpret_cast(&_tmem_load_0[45])), "=r"(*reinterpret_cast(&_tmem_load_0[46])), "=r"(*reinterpret_cast(&_tmem_load_0[47])), "=r"(*reinterpret_cast(&_tmem_load_0[48])), "=r"(*reinterpret_cast(&_tmem_load_0[49])), "=r"(*reinterpret_cast(&_tmem_load_0[50])), "=r"(*reinterpret_cast(&_tmem_load_0[51])), "=r"(*reinterpret_cast(&_tmem_load_0[52])), "=r"(*reinterpret_cast(&_tmem_load_0[53])), "=r"(*reinterpret_cast(&_tmem_load_0[54])), "=r"(*reinterpret_cast(&_tmem_load_0[55])), "=r"(*reinterpret_cast(&_tmem_load_0[56])), "=r"(*reinterpret_cast(&_tmem_load_0[57])), "=r"(*reinterpret_cast(&_tmem_load_0[58])), "=r"(*reinterpret_cast(&_tmem_load_0[59])), "=r"(*reinterpret_cast(&_tmem_load_0[60])), "=r"(*reinterpret_cast(&_tmem_load_0[61])), "=r"(*reinterpret_cast(&_tmem_load_0[62])), "=r"(*reinterpret_cast(&_tmem_load_0[63])) + : "r"(taddr + (unsigned int)score_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + #pragma unroll + for (int rep = 0; rep < 16.0f; rep++) { + int local_reg = rep * 4; + int col_base = score_base + rep * 8 + lane_pair * 2; + float csq0 = smem_csq[col_base]; + float csq1 = smem_csq[col_base + 1]; + float d0 = _tmem_load_0[local_reg] - csq0; + if (d0 > best0) { + best0 = d0; + idx0 = off_k + col_base; + } + float d1 = _tmem_load_0[local_reg + 1] - csq1; + if (d1 > best0) { + best0 = d1; + idx0 = off_k + col_base + 1; + } + float d2 = _tmem_load_0[local_reg + 2] - csq0; + if (d2 > best1) { + best1 = d2; + idx1 = off_k + col_base; + } + float d3 = _tmem_load_0[local_reg + 3] - csq1; + if (d3 > best1) { + best1 = d3; + idx1 = off_k + col_base + 1; + } + } + } + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty_addr); + } + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best0, 1); + float peer0 = _shfl_xor_0; + int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, idx0, 1); + int peer0_idx = _shfl_xor_1; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer0_idx; + } + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, best0, 2); + peer0 = _shfl_xor_2; + int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, idx0, 2); + peer0_idx = _shfl_xor_3; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer0_idx; + } + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, best1, 1); + float peer1 = _shfl_xor_4; + int _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, idx1, 1); + int peer1_idx = _shfl_xor_5; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer1_idx; + } + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, best1, 2); + peer1 = _shfl_xor_6; + int _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, idx1, 2); + peer1_idx = _shfl_xor_7; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer1_idx; + } + if (lane_pair == 0) { + if (global_n0 < N) { + *((int*)(out + out_offset0)) = idx0; + } + if (global_n1 < N) { + *((int*)(out + out_offset1)) = idx1; + } + } + } + } + // ---- Role: load ---- + } else if (warp == 4) { + { // load_main + int num_tiles_1 = B * num_n_tiles; + unsigned int _phase_x_empty_0 = 1; + unsigned int _phase_c_empty_0 = 1; + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int tile_idx_1 = bid; tile_idx_1 < num_tiles_1; tile_idx_1 += num_bids) { + int batch_1 = tile_idx_1 / (unsigned int)num_n_tiles; + int n_tile_1 = tile_idx_1 % (unsigned int)num_n_tiles; + int off_n_1 = n_tile_1 * 64; + int x_row = batch_1 * N + off_n_1; + #pragma unroll 1 + for (int iter_k_1 = 0; iter_k_1 < K_tiles; iter_k_1++) { + int off_k_1 = iter_k_1 * 256; + int c_row = batch_1 * K + off_k_1; + #pragma unroll 1 + for (int iter_d = 0; iter_d < 7; iter_d++) { + int d_group = iter_d; + mbarrier_wait(x_empty_addr, _phase_x_empty_0); + _phase_x_empty_0 ^= 1; + tma_3d_gmem2smem(smem_x_addr, x_tmap, 0, x_row, d_group, x_full_addr); + mbarrier_arrive_expect_tx(x_full_addr, 4096); + mbarrier_wait(c_empty_addr, _phase_c_empty_0); + _phase_c_empty_0 ^= 1; + tma_3d_gmem2smem(smem_c_addr, c_tmap, 0, c_row, d_group, c_full_addr); + mbarrier_arrive_expect_tx(c_full_addr, 16384); + } + } + } + } + } + // ---- Role: mma ---- + } else if (warp == 5) { + { // mma_main + int num_tiles_2 = B * num_n_tiles; + unsigned int _phase_score_empty_0 = 1; + unsigned int _phase_x_full_0 = 0; + unsigned int _phase_c_full_0 = 0; + #pragma unroll 1 + for (unsigned int tile_idx_2 = bid; tile_idx_2 < num_tiles_2; tile_idx_2 += num_bids) { + #pragma unroll 1 + for (int iter_k_2 = 0; iter_k_2 < K_tiles; iter_k_2++) { + mbarrier_wait(score_empty_addr, _phase_score_empty_0); + _phase_score_empty_0 ^= 1; + #pragma unroll 1 + for (int iter_d_1 = 0; iter_d_1 < 7; iter_d_1++) { + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + mbarrier_wait(c_full_addr, _phase_c_full_0); + _phase_c_full_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int init_flag = ((iter_d_1 == 0) ? 1 : 0); + int _mma_a_addr_0 = smem_x_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = smem_c_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 adhi, 0x80004020;\n\t" + "mov.b32 bdhi, 0x80004020;\n\t" + "mov.b32 id, 71304336;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, id, p1;\n\t" + "}\n" + :: "r"(_mma_a_lo_0), "r"(_mma_b_lo_0), "r"(tmem_score_tmem), "r"(((init_flag) ? 0 : 1))); + elect_commit(x_empty_addr); + elect_commit(c_empty_addr); + } + elect_commit(score_full_addr); + } + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0352.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0352.cu new file mode 100644 index 00000000..482ca687 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0352.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 24 * 8; + int row = vec_idx / 24; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 24 * 8; + int row_1 = vec_idx_1 / 24; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0353.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0353.cu new file mode 100644 index 00000000..90ac8670 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0353.cu @@ -0,0 +1,937 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_SX_OFF 1024 +#define SMEM_SX_STAGE_BYTES 14336 +#define SMEM_SX_STRIDE 14336 +#define SMEM_SC0_OFF 15360 +#define SMEM_SC0_STAGE_BYTES 7168 +#define SMEM_SC0_STRIDE 7168 +#define SMEM_SC1_OFF 22528 +#define SMEM_SC1_STAGE_BYTES 7168 +#define SMEM_SC1_STRIDE 7168 +#define SMEM_TOTAL 29696 +#define THREADS 160 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ void cp_async_bulk_gmem2smem( + unsigned smem_addr, const void* gmem_ptr, unsigned bytes, int mbar_addr) { + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(smem_addr), "l"(gmem_ptr), "r"(bytes), "r"(mbar_addr) + : "memory"); +} + +extern "C" { + +__global__ __launch_bounds__(160) void +kernel_flash_kmeans_assign_d112_k1024_owner_local_warp_mma_distinct_workfeed_c829_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* sx = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int sx_addr = smem + 1024; + __nv_bfloat16* sc0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 15360); + const int sc0_addr = smem + 15360; + __nv_bfloat16* sc1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 22528); + const int sc1_addr = smem + 22528; + + // Mbarrier init (5 groups, 5 barriers) + // Mbarriers at smem_raw[0..40) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_full: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // c_full0: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_empty0: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 16, 4, leader); + // c_full1: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_empty1: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 32, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + __syncthreads(); + + const int mbar_base = smem; + #define x_full_addr (mbar_base + 0) + #define c_full0_addr (mbar_base + 8) + #define c_empty0_addr (mbar_base + 16) + #define c_full1_addr (mbar_base + 24) + #define c_empty1_addr (mbar_base + 32) + + // ---- Role: load ---- + if (warp == 0) { + { // load_main + int tile = bid; + int batch = tile / num_n_tiles; + int nt = tile % num_n_tiles; + int point_base = batch * N + nt * 64; + unsigned int _phase_c_empty0_0 = 1; + unsigned int _phase_c_empty1_0 = 1; + if (elect_sync()) { + mbarrier_arrive_expect_tx(x_full_addr, 14336); + cp_async_bulk_gmem2smem(sx_addr, reinterpret_cast(reinterpret_cast(x) + ((unsigned long long)(point_base * D) * (unsigned long long)2)), 14336, x_full_addr); + #pragma unroll 1 + for (unsigned int pair = 0; pair < 16; pair++) { + int kbase0 = pair * 64; + mbarrier_wait(c_empty0_addr, _phase_c_empty0_0); + _phase_c_empty0_0 ^= 1; + mbarrier_arrive_expect_tx(c_full0_addr, 7168); + cp_async_bulk_gmem2smem(sc0_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + kbase0) * D) * (unsigned long long)2)), 7168, c_full0_addr); + int kbase1 = kbase0 + 32; + mbarrier_wait(c_empty1_addr, _phase_c_empty1_0); + _phase_c_empty1_0 ^= 1; + mbarrier_arrive_expect_tx(c_full1_addr, 7168); + cp_async_bulk_gmem2smem(sc1_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + kbase1) * D) * (unsigned long long)2)), 7168, c_full1_addr); + } + } + } + // ---- Role: compute ---- + } else if (warp >= 1 && warp <= 4) { + { // compute_main + int tile_1 = bid; + int batch_1 = tile_1 / num_n_tiles; + int nt_1 = tile_1 % num_n_tiles; + int warp_id_in_role = (warp - 1); + int row_base = warp_id_in_role * 16; + int lane_pair = lane % 4; + float best0 = -3.4e+38f; + float best1 = -3.4e+38f; + int idx0 = 0; + int idx1 = 0; + unsigned int _phase_x_full_0 = 0; + mbarrier_wait(x_full_addr, _phase_x_full_0); + _phase_x_full_0 ^= 1; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + unsigned int _phase_c_full0_0 = 0; + unsigned int _phase_c_full1_0 = 0; + #pragma unroll 1 + for (unsigned int pair_1 = 0; pair_1 < 16; pair_1++) { + int kbase0_1 = pair_1 * 64; + mbarrier_wait(c_full0_addr, _phase_c_full0_0); + _phase_c_full0_0 ^= 1; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + unsigned int a0[4]; + unsigned int b00[2]; + unsigned int b01[2]; + unsigned int b02[2]; + unsigned int b03[2]; + float acc00[4]; + float acc01[4]; + float acc02[4]; + float acc03[4]; + unsigned int a0_addr = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + lane / 16 * 16)); + unsigned int b00_addr = (sc0_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)); + unsigned int b01_addr = (sc0_addr + (unsigned int)((8 + lane % 8) * 224 + (lane / 8 & 1) * 8 * 2)); + unsigned int b02_addr = (sc0_addr + (unsigned int)((16 + lane % 8) * 224 + (lane / 8 & 1) * 8 * 2)); + unsigned int b03_addr = (sc0_addr + (unsigned int)((24 + lane % 8) * 224 + (lane / 8 & 1) * 8 * 2)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a0[0]), "=r"(a0[1]), "=r"(a0[2]), "=r"(a0[3]) + : "r"(a0_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b00[0]), "=r"(b00[1]) + : "r"(b00_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b01[0]), "=r"(b01[1]) + : "r"(b01_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b02[0]), "=r"(b02[1]) + : "r"(b02_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b03[0]), "=r"(b03[1]) + : "r"(b03_addr) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc00[0]), "=f"(acc00[1]), "=f"(acc00[2]), "=f"(acc00[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b00[0]), "r"(b00[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc01[0]), "=f"(acc01[1]), "=f"(acc01[2]), "=f"(acc01[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b01[0]), "r"(b01[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc02[0]), "=f"(acc02[1]), "=f"(acc02[2]), "=f"(acc02[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b02[0]), "r"(b02[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc03[0]), "=f"(acc03[1]), "=f"(acc03[2]), "=f"(acc03[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b03[0]), "r"(b03[1])); + unsigned int a0_addr_0 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 32))); + unsigned int b00_addr_1 = (sc0_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)); + unsigned int b01_addr_2 = (sc0_addr + (unsigned int)((8 + lane % 8) * 224 + (16 + (lane / 8 & 1) * 8) * 2)); + unsigned int b02_addr_3 = (sc0_addr + (unsigned int)((16 + lane % 8) * 224 + (16 + (lane / 8 & 1) * 8) * 2)); + unsigned int b03_addr_4 = (sc0_addr + (unsigned int)((24 + lane % 8) * 224 + (16 + (lane / 8 & 1) * 8) * 2)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a0[0]), "=r"(a0[1]), "=r"(a0[2]), "=r"(a0[3]) + : "r"(a0_addr_0) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b00[0]), "=r"(b00[1]) + : "r"(b00_addr_1) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b01[0]), "=r"(b01[1]) + : "r"(b01_addr_2) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b02[0]), "=r"(b02[1]) + : "r"(b02_addr_3) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b03[0]), "=r"(b03[1]) + : "r"(b03_addr_4) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc00[0]), "+f"(acc00[1]), "+f"(acc00[2]), "+f"(acc00[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b00[0]), "r"(b00[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc01[0]), "+f"(acc01[1]), "+f"(acc01[2]), "+f"(acc01[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b01[0]), "r"(b01[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc02[0]), "+f"(acc02[1]), "+f"(acc02[2]), "+f"(acc02[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b02[0]), "r"(b02[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc03[0]), "+f"(acc03[1]), "+f"(acc03[2]), "+f"(acc03[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b03[0]), "r"(b03[1])); + unsigned int a0_addr_5 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 64))); + unsigned int b00_addr_6 = (sc0_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)); + unsigned int b01_addr_7 = (sc0_addr + (unsigned int)((8 + lane % 8) * 224 + (32 + (lane / 8 & 1) * 8) * 2)); + unsigned int b02_addr_8 = (sc0_addr + (unsigned int)((16 + lane % 8) * 224 + (32 + (lane / 8 & 1) * 8) * 2)); + unsigned int b03_addr_9 = (sc0_addr + (unsigned int)((24 + lane % 8) * 224 + (32 + (lane / 8 & 1) * 8) * 2)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a0[0]), "=r"(a0[1]), "=r"(a0[2]), "=r"(a0[3]) + : "r"(a0_addr_5) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b00[0]), "=r"(b00[1]) + : "r"(b00_addr_6) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b01[0]), "=r"(b01[1]) + : "r"(b01_addr_7) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b02[0]), "=r"(b02[1]) + : "r"(b02_addr_8) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b03[0]), "=r"(b03[1]) + : "r"(b03_addr_9) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc00[0]), "+f"(acc00[1]), "+f"(acc00[2]), "+f"(acc00[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b00[0]), "r"(b00[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc01[0]), "+f"(acc01[1]), "+f"(acc01[2]), "+f"(acc01[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b01[0]), "r"(b01[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc02[0]), "+f"(acc02[1]), "+f"(acc02[2]), "+f"(acc02[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b02[0]), "r"(b02[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc03[0]), "+f"(acc03[1]), "+f"(acc03[2]), "+f"(acc03[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b03[0]), "r"(b03[1])); + unsigned int a0_addr_10 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 96))); + unsigned int b00_addr_11 = (sc0_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)); + unsigned int b01_addr_12 = (sc0_addr + (unsigned int)((8 + lane % 8) * 224 + (48 + (lane / 8 & 1) * 8) * 2)); + unsigned int b02_addr_13 = (sc0_addr + (unsigned int)((16 + lane % 8) * 224 + (48 + (lane / 8 & 1) * 8) * 2)); + unsigned int b03_addr_14 = (sc0_addr + (unsigned int)((24 + lane % 8) * 224 + (48 + (lane / 8 & 1) * 8) * 2)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a0[0]), "=r"(a0[1]), "=r"(a0[2]), "=r"(a0[3]) + : "r"(a0_addr_10) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b00[0]), "=r"(b00[1]) + : "r"(b00_addr_11) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b01[0]), "=r"(b01[1]) + : "r"(b01_addr_12) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b02[0]), "=r"(b02[1]) + : "r"(b02_addr_13) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b03[0]), "=r"(b03[1]) + : "r"(b03_addr_14) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc00[0]), "+f"(acc00[1]), "+f"(acc00[2]), "+f"(acc00[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b00[0]), "r"(b00[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc01[0]), "+f"(acc01[1]), "+f"(acc01[2]), "+f"(acc01[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b01[0]), "r"(b01[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc02[0]), "+f"(acc02[1]), "+f"(acc02[2]), "+f"(acc02[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b02[0]), "r"(b02[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc03[0]), "+f"(acc03[1]), "+f"(acc03[2]), "+f"(acc03[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b03[0]), "r"(b03[1])); + unsigned int a0_addr_15 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 128))); + unsigned int b00_addr_16 = (sc0_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)); + unsigned int b01_addr_17 = (sc0_addr + (unsigned int)((8 + lane % 8) * 224 + (64 + (lane / 8 & 1) * 8) * 2)); + unsigned int b02_addr_18 = (sc0_addr + (unsigned int)((16 + lane % 8) * 224 + (64 + (lane / 8 & 1) * 8) * 2)); + unsigned int b03_addr_19 = (sc0_addr + (unsigned int)((24 + lane % 8) * 224 + (64 + (lane / 8 & 1) * 8) * 2)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a0[0]), "=r"(a0[1]), "=r"(a0[2]), "=r"(a0[3]) + : "r"(a0_addr_15) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b00[0]), "=r"(b00[1]) + : "r"(b00_addr_16) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b01[0]), "=r"(b01[1]) + : "r"(b01_addr_17) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b02[0]), "=r"(b02[1]) + : "r"(b02_addr_18) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b03[0]), "=r"(b03[1]) + : "r"(b03_addr_19) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc00[0]), "+f"(acc00[1]), "+f"(acc00[2]), "+f"(acc00[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b00[0]), "r"(b00[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc01[0]), "+f"(acc01[1]), "+f"(acc01[2]), "+f"(acc01[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b01[0]), "r"(b01[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc02[0]), "+f"(acc02[1]), "+f"(acc02[2]), "+f"(acc02[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b02[0]), "r"(b02[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc03[0]), "+f"(acc03[1]), "+f"(acc03[2]), "+f"(acc03[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b03[0]), "r"(b03[1])); + unsigned int a0_addr_20 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 160))); + unsigned int b00_addr_21 = (sc0_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)); + unsigned int b01_addr_22 = (sc0_addr + (unsigned int)((8 + lane % 8) * 224 + (80 + (lane / 8 & 1) * 8) * 2)); + unsigned int b02_addr_23 = (sc0_addr + (unsigned int)((16 + lane % 8) * 224 + (80 + (lane / 8 & 1) * 8) * 2)); + unsigned int b03_addr_24 = (sc0_addr + (unsigned int)((24 + lane % 8) * 224 + (80 + (lane / 8 & 1) * 8) * 2)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a0[0]), "=r"(a0[1]), "=r"(a0[2]), "=r"(a0[3]) + : "r"(a0_addr_20) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b00[0]), "=r"(b00[1]) + : "r"(b00_addr_21) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b01[0]), "=r"(b01[1]) + : "r"(b01_addr_22) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b02[0]), "=r"(b02[1]) + : "r"(b02_addr_23) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b03[0]), "=r"(b03[1]) + : "r"(b03_addr_24) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc00[0]), "+f"(acc00[1]), "+f"(acc00[2]), "+f"(acc00[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b00[0]), "r"(b00[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc01[0]), "+f"(acc01[1]), "+f"(acc01[2]), "+f"(acc01[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b01[0]), "r"(b01[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc02[0]), "+f"(acc02[1]), "+f"(acc02[2]), "+f"(acc02[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b02[0]), "r"(b02[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc03[0]), "+f"(acc03[1]), "+f"(acc03[2]), "+f"(acc03[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b03[0]), "r"(b03[1])); + unsigned int a0_addr_25 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 192))); + unsigned int b00_addr_26 = (sc0_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)); + unsigned int b01_addr_27 = (sc0_addr + (unsigned int)((8 + lane % 8) * 224 + (96 + (lane / 8 & 1) * 8) * 2)); + unsigned int b02_addr_28 = (sc0_addr + (unsigned int)((16 + lane % 8) * 224 + (96 + (lane / 8 & 1) * 8) * 2)); + unsigned int b03_addr_29 = (sc0_addr + (unsigned int)((24 + lane % 8) * 224 + (96 + (lane / 8 & 1) * 8) * 2)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a0[0]), "=r"(a0[1]), "=r"(a0[2]), "=r"(a0[3]) + : "r"(a0_addr_25) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b00[0]), "=r"(b00[1]) + : "r"(b00_addr_26) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b01[0]), "=r"(b01[1]) + : "r"(b01_addr_27) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b02[0]), "=r"(b02[1]) + : "r"(b02_addr_28) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b03[0]), "=r"(b03[1]) + : "r"(b03_addr_29) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc00[0]), "+f"(acc00[1]), "+f"(acc00[2]), "+f"(acc00[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b00[0]), "r"(b00[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc01[0]), "+f"(acc01[1]), "+f"(acc01[2]), "+f"(acc01[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b01[0]), "r"(b01[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc02[0]), "+f"(acc02[1]), "+f"(acc02[2]), "+f"(acc02[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b02[0]), "r"(b02[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc03[0]), "+f"(acc03[1]), "+f"(acc03[2]), "+f"(acc03[3]) + : "r"(a0[0]), "r"(a0[1]), "r"(a0[2]), "r"(a0[3]), "r"(b03[0]), "r"(b03[1])); + float vals0[8]; + float vals1[8]; + vals0[0] = acc00[0]; + vals0[1] = acc00[1]; + vals0[2] = acc01[0]; + vals0[3] = acc01[1]; + vals0[4] = acc02[0]; + vals0[5] = acc02[1]; + vals0[6] = acc03[0]; + vals0[7] = acc03[1]; + vals1[0] = acc00[2]; + vals1[1] = acc00[3]; + vals1[2] = acc01[2]; + vals1[3] = acc01[3]; + vals1[4] = acc02[2]; + vals1[5] = acc02[3]; + vals1[6] = acc03[2]; + vals1[7] = acc03[3]; + #pragma unroll + for (int vi0 = 0; vi0 < 8; vi0++) { + int col0 = vi0 / 2 * 8 + lane_pair * 2 + vi0 % 2; + float cand0 = vals0[vi0] - 0.5f * c_sq[batch_1 * K + kbase0_1 + col0]; + int cand_idx0 = kbase0_1 + col0; + bool take0 = cand0 > best0; + if (cand0 == best0) { + if (cand_idx0 < idx0) { + take0 = 1; + } + } + if (take0) { + best0 = cand0; + idx0 = cand_idx0; + } + float cand1 = vals1[vi0] - 0.5f * c_sq[batch_1 * K + kbase0_1 + col0]; + int cand_idx1 = kbase0_1 + col0; + bool take1 = cand1 > best1; + if (cand1 == best1) { + if (cand_idx1 < idx1) { + take1 = 1; + } + } + if (take1) { + best1 = cand1; + idx1 = cand_idx1; + } + } + __syncwarp(); + if (elect_sync()) { + mbarrier_arrive(c_empty0_addr); + } + int kbase1_1 = kbase0_1 + 32; + mbarrier_wait(c_full1_addr, _phase_c_full1_0); + _phase_c_full1_0 ^= 1; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + unsigned int a1[4]; + unsigned int b10[2]; + unsigned int b11[2]; + unsigned int b12[2]; + unsigned int b13[2]; + float acc10[4]; + float acc11[4]; + float acc12[4]; + float acc13[4]; + unsigned int a1_addr = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + lane / 16 * 16)); + unsigned int b10_addr = (sc1_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)); + unsigned int b11_addr = (sc1_addr + (unsigned int)((8 + lane % 8) * 224 + (lane / 8 & 1) * 8 * 2)); + unsigned int b12_addr = (sc1_addr + (unsigned int)((16 + lane % 8) * 224 + (lane / 8 & 1) * 8 * 2)); + unsigned int b13_addr = (sc1_addr + (unsigned int)((24 + lane % 8) * 224 + (lane / 8 & 1) * 8 * 2)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a1[0]), "=r"(a1[1]), "=r"(a1[2]), "=r"(a1[3]) + : "r"(a1_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b10[0]), "=r"(b10[1]) + : "r"(b10_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b11[0]), "=r"(b11[1]) + : "r"(b11_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b12[0]), "=r"(b12[1]) + : "r"(b12_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b13[0]), "=r"(b13[1]) + : "r"(b13_addr) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc10[0]), "=f"(acc10[1]), "=f"(acc10[2]), "=f"(acc10[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b10[0]), "r"(b10[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc11[0]), "=f"(acc11[1]), "=f"(acc11[2]), "=f"(acc11[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b11[0]), "r"(b11[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc12[0]), "=f"(acc12[1]), "=f"(acc12[2]), "=f"(acc12[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b12[0]), "r"(b12[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc13[0]), "=f"(acc13[1]), "=f"(acc13[2]), "=f"(acc13[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b13[0]), "r"(b13[1])); + unsigned int a1_addr_30 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 32))); + unsigned int b10_addr_31 = (sc1_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)); + unsigned int b11_addr_32 = (sc1_addr + (unsigned int)((8 + lane % 8) * 224 + (16 + (lane / 8 & 1) * 8) * 2)); + unsigned int b12_addr_33 = (sc1_addr + (unsigned int)((16 + lane % 8) * 224 + (16 + (lane / 8 & 1) * 8) * 2)); + unsigned int b13_addr_34 = (sc1_addr + (unsigned int)((24 + lane % 8) * 224 + (16 + (lane / 8 & 1) * 8) * 2)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a1[0]), "=r"(a1[1]), "=r"(a1[2]), "=r"(a1[3]) + : "r"(a1_addr_30) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b10[0]), "=r"(b10[1]) + : "r"(b10_addr_31) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b11[0]), "=r"(b11[1]) + : "r"(b11_addr_32) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b12[0]), "=r"(b12[1]) + : "r"(b12_addr_33) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b13[0]), "=r"(b13[1]) + : "r"(b13_addr_34) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc10[0]), "+f"(acc10[1]), "+f"(acc10[2]), "+f"(acc10[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b10[0]), "r"(b10[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc11[0]), "+f"(acc11[1]), "+f"(acc11[2]), "+f"(acc11[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b11[0]), "r"(b11[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc12[0]), "+f"(acc12[1]), "+f"(acc12[2]), "+f"(acc12[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b12[0]), "r"(b12[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc13[0]), "+f"(acc13[1]), "+f"(acc13[2]), "+f"(acc13[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b13[0]), "r"(b13[1])); + unsigned int a1_addr_35 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 64))); + unsigned int b10_addr_36 = (sc1_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)); + unsigned int b11_addr_37 = (sc1_addr + (unsigned int)((8 + lane % 8) * 224 + (32 + (lane / 8 & 1) * 8) * 2)); + unsigned int b12_addr_38 = (sc1_addr + (unsigned int)((16 + lane % 8) * 224 + (32 + (lane / 8 & 1) * 8) * 2)); + unsigned int b13_addr_39 = (sc1_addr + (unsigned int)((24 + lane % 8) * 224 + (32 + (lane / 8 & 1) * 8) * 2)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a1[0]), "=r"(a1[1]), "=r"(a1[2]), "=r"(a1[3]) + : "r"(a1_addr_35) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b10[0]), "=r"(b10[1]) + : "r"(b10_addr_36) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b11[0]), "=r"(b11[1]) + : "r"(b11_addr_37) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b12[0]), "=r"(b12[1]) + : "r"(b12_addr_38) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b13[0]), "=r"(b13[1]) + : "r"(b13_addr_39) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc10[0]), "+f"(acc10[1]), "+f"(acc10[2]), "+f"(acc10[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b10[0]), "r"(b10[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc11[0]), "+f"(acc11[1]), "+f"(acc11[2]), "+f"(acc11[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b11[0]), "r"(b11[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc12[0]), "+f"(acc12[1]), "+f"(acc12[2]), "+f"(acc12[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b12[0]), "r"(b12[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc13[0]), "+f"(acc13[1]), "+f"(acc13[2]), "+f"(acc13[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b13[0]), "r"(b13[1])); + unsigned int a1_addr_40 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 96))); + unsigned int b10_addr_41 = (sc1_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)); + unsigned int b11_addr_42 = (sc1_addr + (unsigned int)((8 + lane % 8) * 224 + (48 + (lane / 8 & 1) * 8) * 2)); + unsigned int b12_addr_43 = (sc1_addr + (unsigned int)((16 + lane % 8) * 224 + (48 + (lane / 8 & 1) * 8) * 2)); + unsigned int b13_addr_44 = (sc1_addr + (unsigned int)((24 + lane % 8) * 224 + (48 + (lane / 8 & 1) * 8) * 2)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a1[0]), "=r"(a1[1]), "=r"(a1[2]), "=r"(a1[3]) + : "r"(a1_addr_40) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b10[0]), "=r"(b10[1]) + : "r"(b10_addr_41) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b11[0]), "=r"(b11[1]) + : "r"(b11_addr_42) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b12[0]), "=r"(b12[1]) + : "r"(b12_addr_43) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b13[0]), "=r"(b13[1]) + : "r"(b13_addr_44) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc10[0]), "+f"(acc10[1]), "+f"(acc10[2]), "+f"(acc10[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b10[0]), "r"(b10[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc11[0]), "+f"(acc11[1]), "+f"(acc11[2]), "+f"(acc11[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b11[0]), "r"(b11[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc12[0]), "+f"(acc12[1]), "+f"(acc12[2]), "+f"(acc12[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b12[0]), "r"(b12[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc13[0]), "+f"(acc13[1]), "+f"(acc13[2]), "+f"(acc13[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b13[0]), "r"(b13[1])); + unsigned int a1_addr_45 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 128))); + unsigned int b10_addr_46 = (sc1_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)); + unsigned int b11_addr_47 = (sc1_addr + (unsigned int)((8 + lane % 8) * 224 + (64 + (lane / 8 & 1) * 8) * 2)); + unsigned int b12_addr_48 = (sc1_addr + (unsigned int)((16 + lane % 8) * 224 + (64 + (lane / 8 & 1) * 8) * 2)); + unsigned int b13_addr_49 = (sc1_addr + (unsigned int)((24 + lane % 8) * 224 + (64 + (lane / 8 & 1) * 8) * 2)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a1[0]), "=r"(a1[1]), "=r"(a1[2]), "=r"(a1[3]) + : "r"(a1_addr_45) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b10[0]), "=r"(b10[1]) + : "r"(b10_addr_46) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b11[0]), "=r"(b11[1]) + : "r"(b11_addr_47) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b12[0]), "=r"(b12[1]) + : "r"(b12_addr_48) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b13[0]), "=r"(b13[1]) + : "r"(b13_addr_49) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc10[0]), "+f"(acc10[1]), "+f"(acc10[2]), "+f"(acc10[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b10[0]), "r"(b10[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc11[0]), "+f"(acc11[1]), "+f"(acc11[2]), "+f"(acc11[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b11[0]), "r"(b11[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc12[0]), "+f"(acc12[1]), "+f"(acc12[2]), "+f"(acc12[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b12[0]), "r"(b12[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc13[0]), "+f"(acc13[1]), "+f"(acc13[2]), "+f"(acc13[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b13[0]), "r"(b13[1])); + unsigned int a1_addr_50 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 160))); + unsigned int b10_addr_51 = (sc1_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)); + unsigned int b11_addr_52 = (sc1_addr + (unsigned int)((8 + lane % 8) * 224 + (80 + (lane / 8 & 1) * 8) * 2)); + unsigned int b12_addr_53 = (sc1_addr + (unsigned int)((16 + lane % 8) * 224 + (80 + (lane / 8 & 1) * 8) * 2)); + unsigned int b13_addr_54 = (sc1_addr + (unsigned int)((24 + lane % 8) * 224 + (80 + (lane / 8 & 1) * 8) * 2)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a1[0]), "=r"(a1[1]), "=r"(a1[2]), "=r"(a1[3]) + : "r"(a1_addr_50) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b10[0]), "=r"(b10[1]) + : "r"(b10_addr_51) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b11[0]), "=r"(b11[1]) + : "r"(b11_addr_52) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b12[0]), "=r"(b12[1]) + : "r"(b12_addr_53) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b13[0]), "=r"(b13[1]) + : "r"(b13_addr_54) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc10[0]), "+f"(acc10[1]), "+f"(acc10[2]), "+f"(acc10[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b10[0]), "r"(b10[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc11[0]), "+f"(acc11[1]), "+f"(acc11[2]), "+f"(acc11[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b11[0]), "r"(b11[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc12[0]), "+f"(acc12[1]), "+f"(acc12[2]), "+f"(acc12[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b12[0]), "r"(b12[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc13[0]), "+f"(acc13[1]), "+f"(acc13[2]), "+f"(acc13[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b13[0]), "r"(b13[1])); + unsigned int a1_addr_55 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 192))); + unsigned int b10_addr_56 = (sc1_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)); + unsigned int b11_addr_57 = (sc1_addr + (unsigned int)((8 + lane % 8) * 224 + (96 + (lane / 8 & 1) * 8) * 2)); + unsigned int b12_addr_58 = (sc1_addr + (unsigned int)((16 + lane % 8) * 224 + (96 + (lane / 8 & 1) * 8) * 2)); + unsigned int b13_addr_59 = (sc1_addr + (unsigned int)((24 + lane % 8) * 224 + (96 + (lane / 8 & 1) * 8) * 2)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a1[0]), "=r"(a1[1]), "=r"(a1[2]), "=r"(a1[3]) + : "r"(a1_addr_55) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b10[0]), "=r"(b10[1]) + : "r"(b10_addr_56) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b11[0]), "=r"(b11[1]) + : "r"(b11_addr_57) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b12[0]), "=r"(b12[1]) + : "r"(b12_addr_58) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b13[0]), "=r"(b13[1]) + : "r"(b13_addr_59) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc10[0]), "+f"(acc10[1]), "+f"(acc10[2]), "+f"(acc10[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b10[0]), "r"(b10[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc11[0]), "+f"(acc11[1]), "+f"(acc11[2]), "+f"(acc11[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b11[0]), "r"(b11[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc12[0]), "+f"(acc12[1]), "+f"(acc12[2]), "+f"(acc12[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b12[0]), "r"(b12[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc13[0]), "+f"(acc13[1]), "+f"(acc13[2]), "+f"(acc13[3]) + : "r"(a1[0]), "r"(a1[1]), "r"(a1[2]), "r"(a1[3]), "r"(b13[0]), "r"(b13[1])); + float vals10[8]; + float vals11[8]; + vals10[0] = acc10[0]; + vals10[1] = acc10[1]; + vals10[2] = acc11[0]; + vals10[3] = acc11[1]; + vals10[4] = acc12[0]; + vals10[5] = acc12[1]; + vals10[6] = acc13[0]; + vals10[7] = acc13[1]; + vals11[0] = acc10[2]; + vals11[1] = acc10[3]; + vals11[2] = acc11[2]; + vals11[3] = acc11[3]; + vals11[4] = acc12[2]; + vals11[5] = acc12[3]; + vals11[6] = acc13[2]; + vals11[7] = acc13[3]; + #pragma unroll + for (int vi1 = 0; vi1 < 8; vi1++) { + int col1 = vi1 / 2 * 8 + lane_pair * 2 + vi1 % 2; + float cand10 = vals10[vi1] - 0.5f * c_sq[batch_1 * K + kbase1_1 + col1]; + int cand_idx10 = kbase1_1 + col1; + bool take10 = cand10 > best0; + if (cand10 == best0) { + if (cand_idx10 < idx0) { + take10 = 1; + } + } + if (take10) { + best0 = cand10; + idx0 = cand_idx10; + } + float cand11 = vals11[vi1] - 0.5f * c_sq[batch_1 * K + kbase1_1 + col1]; + int cand_idx11 = kbase1_1 + col1; + bool take11 = cand11 > best1; + if (cand11 == best1) { + if (cand_idx11 < idx1) { + take11 = 1; + } + } + if (take11) { + best1 = cand11; + idx1 = cand_idx11; + } + } + __syncwarp(); + if (elect_sync()) { + mbarrier_arrive(c_empty1_addr); + } + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best0, 1); + float peer0 = _shfl_xor_0; + int _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, idx0, 1); + int peer_idx0 = _shfl_xor_1; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer_idx0; + } + if (peer0 == best0) { + if (peer_idx0 < idx0) { + idx0 = peer_idx0; + } + } + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, best0, 2); + peer0 = _shfl_xor_2; + int _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, idx0, 2); + peer_idx0 = _shfl_xor_3; + if (peer0 > best0) { + best0 = peer0; + idx0 = peer_idx0; + } + if (peer0 == best0) { + if (peer_idx0 < idx0) { + idx0 = peer_idx0; + } + } + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, best1, 1); + float peer1 = _shfl_xor_4; + int _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, idx1, 1); + int peer_idx1 = _shfl_xor_5; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer_idx1; + } + if (peer1 == best1) { + if (peer_idx1 < idx1) { + idx1 = peer_idx1; + } + } + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, best1, 2); + peer1 = _shfl_xor_6; + int _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, idx1, 2); + peer_idx1 = _shfl_xor_7; + if (peer1 > best1) { + best1 = peer1; + idx1 = peer_idx1; + } + if (peer1 == best1) { + if (peer_idx1 < idx1) { + idx1 = peer_idx1; + } + } + if (lane_pair == 0) { + int row0 = row_base + lane / 4; + *((int*)(out + (batch_1 * N + nt_1 * 64 + row0))) = idx0; + *((int*)(out + (batch_1 * N + nt_1 * 64 + row0 + 8))) = idx1; + } + } + } + + // Cleanup +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0354.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0354.cu new file mode 100644 index 00000000..4f6e1450 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0354.cu @@ -0,0 +1,540 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_X_RAW_OFF 1024 +#define SMEM_X_RAW_STAGE_BYTES 14336 +#define SMEM_X_RAW_STRIDE 14336 +#define SMEM_C_DIRECT00_OFF 15360 +#define SMEM_C_DIRECT00_STAGE_BYTES 1792 +#define SMEM_C_DIRECT00_STRIDE 1792 +#define SMEM_C_DIRECT01_OFF 17152 +#define SMEM_C_DIRECT01_STAGE_BYTES 1792 +#define SMEM_C_DIRECT01_STRIDE 1792 +#define SMEM_C_DIRECT10_OFF 18944 +#define SMEM_C_DIRECT10_STAGE_BYTES 1792 +#define SMEM_C_DIRECT10_STRIDE 1792 +#define SMEM_C_DIRECT11_OFF 20736 +#define SMEM_C_DIRECT11_STAGE_BYTES 1792 +#define SMEM_C_DIRECT11_STRIDE 1792 +#define SMEM_SX_OFF 22528 +#define SMEM_SX_STAGE_BYTES 14336 +#define SMEM_SX_STRIDE 14336 +#define SMEM_SS_OFF 36864 +#define SMEM_SS_STAGE_BYTES 4096 +#define SMEM_SS_STRIDE 4096 +#define SMEM_GROUP_KEYS_OFF 40960 +#define SMEM_GROUP_KEYS_STAGE_BYTES 1024 +#define SMEM_GROUP_KEYS_STRIDE 1024 +#define SMEM_LOCAL_KEYS_OFF 41984 +#define SMEM_LOCAL_KEYS_STAGE_BYTES 512 +#define SMEM_LOCAL_KEYS_STRIDE 512 +#define SMEM_PEER_KEYS_OFF 42496 +#define SMEM_PEER_KEYS_STAGE_BYTES 512 +#define SMEM_PEER_KEYS_STRIDE 512 +#define SMEM_TOTAL 43008 +#define THREADS 256 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ void cp_async_bulk_gmem2smem( + unsigned smem_addr, const void* gmem_ptr, unsigned bytes, int mbar_addr) { + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(smem_addr), "l"(gmem_ptr), "r"(bytes), "r"(mbar_addr) + : "memory"); +} + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d112_k512_two_owner_peer_only_mma_f826_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* x_raw = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int x_raw_addr = smem + 1024; + __nv_bfloat16* c_direct00 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 15360); + const int c_direct00_addr = smem + 15360; + __nv_bfloat16* c_direct01 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17152); + const int c_direct01_addr = smem + 17152; + __nv_bfloat16* c_direct10 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 18944); + const int c_direct10_addr = smem + 18944; + __nv_bfloat16* c_direct11 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 20736); + const int c_direct11_addr = smem + 20736; + __nv_bfloat16* sx = reinterpret_cast<__nv_bfloat16*>(smem_raw + 22528); + const int sx_addr = smem + 22528; + float* ss = reinterpret_cast(smem_raw + 36864); + const int ss_addr = smem + 36864; + unsigned long long* group_keys = reinterpret_cast(smem_raw + 40960); + const int group_keys_addr = smem + 40960; + unsigned long long* local_keys = reinterpret_cast(smem_raw + 41984); + const int local_keys_addr = smem + 41984; + unsigned long long* peer_keys = reinterpret_cast(smem_raw + 42496); + const int peer_keys_addr = smem + 42496; + + // Mbarrier init (6 groups, 6 barriers) + // Mbarriers at smem_raw[0..48) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_ready: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // c_ready00: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_ready01: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_ready10: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_ready11: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // peer_keys_ready: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); + + const int mbar_base = smem; + #define x_ready_addr (mbar_base + 0) + #define c_ready00_addr (mbar_base + 8) + #define c_ready01_addr (mbar_base + 16) + #define c_ready10_addr (mbar_base + 24) + #define c_ready11_addr (mbar_base + 32) + #define peer_keys_ready_addr (mbar_base + 40) + + // === Task calls (dependency order) === + int total_tiles = B * num_n_tiles; + unsigned int _phase_x_ready_0 = 0; + unsigned int _phase_c_ready00_0 = 0; + unsigned int _phase_c_ready01_0 = 0; + unsigned int _phase_c_ready10_0 = 0; + unsigned int _phase_c_ready11_0 = 0; + unsigned int _phase_peer_keys_ready_0 = 0; + #pragma unroll 1 + for (unsigned int tile = cluster_id; tile < total_tiles; tile += num_clusters) { + int batch = tile / (unsigned int)num_n_tiles; + int nt = tile % (unsigned int)num_n_tiles; + int point_base = batch * N + nt * 64; + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(x_ready_addr, 14336); + cp_async_bulk_gmem2smem(x_raw_addr, reinterpret_cast(reinterpret_cast(x) + ((unsigned long long)(point_base * D) * (unsigned long long)2)), 14336, x_ready_addr); + } + } + mbarrier_wait(x_ready_addr, _phase_x_ready_0); + _phase_x_ready_0 ^= 1; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + #pragma unroll 1 + for (unsigned int i = tid; i < 7168; i += 256) { + int row = i / 112; + int col = i % 112; + { + __nv_bfloat16 _bval_3404602512 = __float2bfloat16_rn(x_raw[i]); + uint16_t _bits_3404602512 = *(uint16_t*)&_bval_3404602512; + uint32_t _addr_3404602512 = static_cast((sx_addr + (unsigned int)(row * 224 + col * 2))); + asm volatile("st.shared.b16 [%0], %1;" :: "r"(_addr_3404602512), "h"(_bits_3404602512) : "memory"); + } + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + int warp_id_in_role = (warp - 0); + int group = warp_id_in_role / 4; + int local_warp = warp_id_in_role % 4; + int local_row = lane % 16; + int row_base = local_warp * 16; + float best = -3.4e+38f; + int owner_k_tiles = K / 16; + int group_k_tiles = owner_k_tiles / 2; + int group_tile_begin = cta_rank * owner_k_tiles + group * group_k_tiles; + int best_idx = group_tile_begin * 8; + int first_kbase = group_tile_begin * 8; + if (group == 0) { + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready00_addr, 1792); + cp_async_bulk_gmem2smem(c_direct00_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + first_kbase) * D) * (unsigned long long)2)), 1792, c_ready00_addr); + } + } + } else if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready10_addr, 1792); + cp_async_bulk_gmem2smem(c_direct10_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + first_kbase) * D) * (unsigned long long)2)), 1792, c_ready10_addr); + } + } + #pragma unroll 1 + for (int group_kt = 0; group_kt < group_k_tiles; group_kt++) { + int current_stage = group_kt % 2; + int kt = group_tile_begin + group_kt; + int kbase = kt * 8; + bool has_next = group_k_tiles > group_kt + 1; + if (has_next) { + int next_stage = (group_kt + 1) % 2; + int next_kbase = (kt + 1) * 8; + if (group == 0) { + if (next_stage == 0) { + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready00_addr, 1792); + cp_async_bulk_gmem2smem(c_direct00_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready00_addr); + } + } + } else if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready01_addr, 1792); + cp_async_bulk_gmem2smem(c_direct01_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready01_addr); + } + } + } else if (next_stage == 0) { + if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready10_addr, 1792); + cp_async_bulk_gmem2smem(c_direct10_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready10_addr); + } + } + } else { + if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready11_addr, 1792); + cp_async_bulk_gmem2smem(c_direct11_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready11_addr); + } + } + } + } + if (group == 0) { + if (current_stage == 0) { + mbarrier_wait(c_ready00_addr, _phase_c_ready00_0); + _phase_c_ready00_0 ^= 1; + } else { + mbarrier_wait(c_ready01_addr, _phase_c_ready01_0); + _phase_c_ready01_0 ^= 1; + } + } else if (current_stage == 0) { + mbarrier_wait(c_ready10_addr, _phase_c_ready10_0); + _phase_c_ready10_0 ^= 1; + } else { + mbarrier_wait(c_ready11_addr, _phase_c_ready11_0); + _phase_c_ready11_0 ^= 1; + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + unsigned int a[4]; + unsigned int b[2]; + float acc[4]; + unsigned int a_addr = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + lane / 16 * 16)); + unsigned int b_addr = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[0]), "=f"(acc[1]), "=f"(acc[2]), "=f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_0 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 32))); + unsigned int b_addr_1 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_0) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_1) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_2 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 64))); + unsigned int b_addr_3 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_2) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_3) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_4 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 96))); + unsigned int b_addr_5 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_4) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_5) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_6 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 128))); + unsigned int b_addr_7 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_6) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_7) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_8 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 160))); + unsigned int b_addr_9 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_8) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_9) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_10 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 192))); + unsigned int b_addr_11 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_10) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_11) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + #pragma unroll + for (int rp = 0; rp < 2; rp++) { + #pragma unroll + for (int cp = 0; cp < 2; cp++) { + int rr = lane / 4 + rp * 8; + int cc = lane % 4 * 2 + cp; + { + uint32_t _addr_3405267776 = static_cast((ss_addr + (unsigned int)((group * 64 + row_base + rr) * 32 + cc * 4))); + asm volatile("st.shared.f32 [%0], %1;" :: "r"(_addr_3405267776), "f"(acc[rp * 2 + cp]) : "memory"); + } + } + } + __syncwarp(); + if (lane < 16) { + #pragma unroll + for (int kk = 0; kk < 8; kk++) { + float score = ss[(group * 64 + row_base + local_row) * 8 + kk] - 0.5f * c_sq[batch * K + kbase + kk]; + if (score > best) { + best = score; + best_idx = kbase + kk; + } + } + } + if (group == 0) { + asm volatile("barrier.sync 1, %0;" :: "r"(128)); + } else { + asm volatile("barrier.sync 2, %0;" :: "r"(128)); + } + } + if (lane < 16) { + uint32_t _amf_u_0 = __float_as_uint(best); + uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; + unsigned int _amf_enc_0 = _amf_u_0 ^ _amf_mask_0; + unsigned long long shift64 = 32; + unsigned long long mask64 = 4294967295; + group_keys[group * 64 + row_base + local_row] = (unsigned long long)_amf_enc_0 << shift64 | mask64 - (unsigned long long)best_idx; + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + if (group == 0 && lane < 16) { + unsigned long long first_key = group_keys[row_base + local_row]; + unsigned long long second_key = group_keys[64 + row_base + local_row]; + local_keys[row_base + local_row] = ((second_key > first_key) ? second_key : first_key); + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + if (cta_rank == 1) { + if (warp == 0) { + if (elect_sync()) { + uint32_t _mapa_0; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_0) : "r"(peer_keys_addr), "r"(0)); + uint32_t _mapa_1; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_1) : "r"(peer_keys_ready_addr), "r"(0)); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile( + "mbarrier.arrive.expect_tx.release.cluster.shared::cluster.b64 _, [%0], %1;" + :: "r"(_mapa_1), "r"((uint32_t)(512)) : "memory"); + asm volatile( + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(_mapa_0), "r"(local_keys_addr), "r"((uint32_t)(512)), "r"(_mapa_1) + : "memory"); + } + } + } + if (cta_rank == 0) { + mbarrier_wait(peer_keys_ready_addr, _phase_peer_keys_ready_0); + _phase_peer_keys_ready_0 ^= 1; + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + int row_1 = tid; + if (row_1 < 64) { + unsigned long long local_key = local_keys[row_1]; + unsigned long long peer_key = peer_keys[row_1]; + unsigned long long best_key = ((peer_key > local_key) ? peer_key : local_key); + unsigned long long mask64_1 = 4294967295; + int idx = (int)(mask64_1 - (best_key & mask64_1)); + *((int*)(out + (batch * N + nt * 64 + row_1))) = idx; + } + } + } + + // Cleanup + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0355.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0355.cu new file mode 100644 index 00000000..b875687d --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0355.cu @@ -0,0 +1,628 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 128 +#define TMEM_SCORES0_OFFSET 0 +#define TMEM_SCORES1_OFFSET 64 +#define NUM_X_PIPE_STAGES 7 +#define SMEM_SX_OFF 1024 +#define SMEM_SX_STAGE_BYTES 4096 +#define SMEM_SX_STRIDE 4096 +#define SMEM_SC0_OFF 29696 +#define SMEM_SC0_STAGE_BYTES 2048 +#define SMEM_SC0_STRIDE 2048 +#define SMEM_SC1_OFF 31744 +#define SMEM_SC1_STAGE_BYTES 2048 +#define SMEM_SC1_STRIDE 2048 +#define SMEM_SLOT0_KEYS_OFF 33792 +#define SMEM_SLOT0_KEYS_STAGE_BYTES 1024 +#define SMEM_SLOT0_KEYS_STRIDE 1024 +#define SMEM_SLOT1_KEYS_OFF 34816 +#define SMEM_SLOT1_KEYS_STAGE_BYTES 1024 +#define SMEM_SLOT1_KEYS_STRIDE 1024 +#define SMEM_TOTAL 35840 +#define THREADS 416 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ss_step( + int a_lo, int b_lo, int taddr, uint32_t i_desc, int enable_d, + uint32_t a_dhi, uint32_t b_dhi) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 adhi, bdhi;\n\t" + ".reg .b64 da, db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "mov.b32 adhi, %5;\n\t" + "mov.b32 bdhi, %6;\n\t" + "mov.b64 da, {%0, adhi};\n\t" + "mov.b64 db, {%1, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%2], da, db, %3, p;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"(taddr), "r"(i_desc), "r"(enable_d), "r"(a_dhi), "r"(b_dhi)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16_wait(float* dst, int addr) { + tmem_ld_x16(dst, addr); + asm volatile("tcgen05.wait::ld.sync.aligned;"); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(416) void +kernel_flash_kmeans_assign_d112_shared_point_dual_issuer_tcgen05_6e1e_v1(const void* __restrict__ x_tmap, const void* __restrict__ c_tmap, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles, int K_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* sx = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int sx_addr = smem + 1024; + __nv_bfloat16* sc0 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 29696); + const int sc0_addr = smem + 29696; + __nv_bfloat16* sc1 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 31744); + const int sc1_addr = smem + 31744; + unsigned long long* slot0_keys = reinterpret_cast(smem_raw + 33792); + const int slot0_keys_addr = smem + 33792; + unsigned long long* slot1_keys = reinterpret_cast(smem_raw + 34816); + const int slot1_keys_addr = smem + 34816; + + // Mbarrier init (14 groups, 20 barriers) + // Mbarriers at smem_raw[0..160) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // --- pipeline 'x_pipe' --- + // x_tma_full: 7 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + mbarrier_init_pred(smem + 8, 1, leader); + mbarrier_init_pred(smem + 16, 1, leader); + mbarrier_init_pred(smem + 24, 1, leader); + mbarrier_init_pred(smem + 32, 1, leader); + mbarrier_init_pred(smem + 40, 1, leader); + mbarrier_init_pred(smem + 48, 1, leader); + // x_ready0: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 56, 1, leader); + // x_empty0: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 64, 1, leader); + // x_ready1: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 72, 1, leader); + // x_empty1: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 80, 1, leader); + // c_full0: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 88, 1, leader); + // c_empty0: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 96, 1, leader); + // c_full1: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 104, 1, leader); + // c_empty1: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 112, 1, leader); + // score_full0: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 120, 1, leader); + // score_empty0: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 128, 4, leader); + // score_full1: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 136, 1, leader); + // score_empty1: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 144, 4, leader); + // slot1_keys_ready: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 152, 4, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (128 columns, 128 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 160); + if (warp == 0) { + int _tmem_hold = smem + 160; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(128) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define x_tma_full_addr (mbar_base + 0) + #define x_ready0_addr (mbar_base + 56) + #define x_empty0_addr (mbar_base + 64) + #define x_ready1_addr (mbar_base + 72) + #define x_empty1_addr (mbar_base + 80) + #define c_full0_addr (mbar_base + 88) + #define c_empty0_addr (mbar_base + 96) + #define c_full1_addr (mbar_base + 104) + #define c_empty1_addr (mbar_base + 112) + #define score_full0_addr (mbar_base + 120) + #define score_empty0_addr (mbar_base + 128) + #define score_full1_addr (mbar_base + 136) + #define score_empty1_addr (mbar_base + 144) + #define slot1_keys_ready_addr (mbar_base + 152) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_scores0 = taddr; + const int tmem_scores1 = taddr + 64; + + // ---- Role: x_load ---- + if (warp == 0) { + { // x_load_main + int total_tiles = B * num_n_tiles; + unsigned int _phase_x_empty0_0 = 1; + unsigned int _phase_x_empty1_0 = 1; + unsigned int _phase_x_tma_full = 0; + #pragma unroll 1 + for (unsigned int tile = bid; tile < total_tiles; tile += num_bids) { + int point_row = tile / (unsigned int)num_n_tiles * (unsigned int)N + tile % (unsigned int)num_n_tiles * 128; + mbarrier_wait(x_empty0_addr, _phase_x_empty0_0); + _phase_x_empty0_0 ^= 1; + mbarrier_wait(x_empty1_addr, _phase_x_empty1_0); + _phase_x_empty1_0 ^= 1; + if (warp == 0) { + if (elect_sync()) { + #pragma unroll + for (int ft = 0; ft < 7; ft++) { + mbarrier_arrive_expect_tx(x_tma_full_addr + (ft) * 8, 4096); + tma_3d_gmem2smem(sx_addr + (unsigned int)(ft * 4096), x_tmap, 0, point_row, ft, x_tma_full_addr + (ft) * 8); + } + } + } + #pragma unroll + for (int ft_1 = 0; ft_1 < 7; ft_1++) { + mbarrier_wait(x_tma_full_addr + (ft_1) * 8, _phase_x_tma_full); + } + asm volatile("fence.proxy.async;"); + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive(x_ready0_addr); + mbarrier_arrive(x_ready1_addr); + } + } + } + } + // ---- Role: c_load0 ---- + } else if (warp == 1) { + { // c_load0_main + int total_tiles_1 = B * num_n_tiles; + int owner_tiles = K_tiles / 2; + unsigned int _phase_c_empty0_0 = 1; + #pragma unroll 1 + for (unsigned int tile_1 = bid; tile_1 < total_tiles_1; tile_1 += num_bids) { + int batch = tile_1 / (unsigned int)num_n_tiles; + if (warp == 1) { + if (elect_sync()) { + #pragma unroll 1 + for (int owner_kt = 0; owner_kt < owner_tiles; owner_kt++) { + int kt = owner_kt * 2; + int centroid_row = batch * K + kt * 64; + #pragma unroll + for (int ft_2 = 0; ft_2 < 7; ft_2++) { + mbarrier_wait(c_empty0_addr, _phase_c_empty0_0); + _phase_c_empty0_0 ^= 1; + mbarrier_arrive_expect_tx(c_full0_addr, 2048); + tma_3d_gmem2smem(sc0_addr, c_tmap, 0, centroid_row, ft_2, c_full0_addr); + } + } + } + } + } + } + // ---- Role: c_load1 ---- + } else if (warp == 2) { + { // c_load1_main + int total_tiles_2 = B * num_n_tiles; + int owner_tiles_1 = K_tiles / 2; + unsigned int _phase_c_empty1_0 = 1; + #pragma unroll 1 + for (unsigned int tile_2 = bid; tile_2 < total_tiles_2; tile_2 += num_bids) { + int batch_1 = tile_2 / (unsigned int)num_n_tiles; + if (warp == 2) { + if (elect_sync()) { + #pragma unroll 1 + for (int owner_kt_1 = 0; owner_kt_1 < owner_tiles_1; owner_kt_1++) { + int kt_1 = owner_kt_1 * 2 + 1; + int centroid_row_1 = batch_1 * K + kt_1 * 64; + #pragma unroll + for (int ft_3 = 0; ft_3 < 7; ft_3++) { + mbarrier_wait(c_empty1_addr, _phase_c_empty1_0); + _phase_c_empty1_0 ^= 1; + mbarrier_arrive_expect_tx(c_full1_addr, 2048); + tma_3d_gmem2smem(sc1_addr, c_tmap, 0, centroid_row_1, ft_3, c_full1_addr); + } + } + } + } + } + } + // ---- Role: mma0 ---- + } else if (warp == 3) { + { // mma0_main + int total_tiles_3 = B * num_n_tiles; + int owner_tiles_2 = K_tiles / 2; + unsigned int _phase_x_ready0_0 = 0; + unsigned int _phase_score_empty0_0 = 1; + unsigned int _phase_c_full0_0 = 0; + #pragma unroll 1 + for (unsigned int tile_3 = bid; tile_3 < total_tiles_3; tile_3 += num_bids) { + mbarrier_wait(x_ready0_addr, _phase_x_ready0_0); + _phase_x_ready0_0 ^= 1; + #pragma unroll 1 + for (int owner_kt_2 = 0; owner_kt_2 < owner_tiles_2; owner_kt_2++) { + mbarrier_wait(score_empty0_addr, _phase_score_empty0_0); + _phase_score_empty0_0 ^= 1; + mbarrier_wait(c_full0_addr, _phase_c_full0_0); + _phase_c_full0_0 ^= 1; + asm volatile("fence.proxy.async;"); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_0 = sx_addr; + int _mma_a_lo_0 = make_warp_uniform((_mma_a_addr_0 >> 4) & 0x3FFF); + int _mma_b_addr_0 = sc0_addr; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + mma_ss_step(_mma_a_lo_0, _mma_b_lo_0, tmem_scores0, 135267472, 0, 0xC0004010U, 0xC0004010U); + elect_commit(c_empty0_addr); + #pragma unroll + for (int ft_4 = 1; ft_4 < 7; ft_4++) { + mbarrier_wait(c_full0_addr, _phase_c_full0_0); + _phase_c_full0_0 ^= 1; + asm volatile("fence.proxy.async;"); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_1 = sx_addr + (unsigned int)(ft_4 * 4096); + int _mma_a_lo_1 = make_warp_uniform((_mma_a_addr_1 >> 4) & 0x3FFF); + int _mma_b_addr_1 = sc0_addr; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + mma_ss_step(_mma_a_lo_1, _mma_b_lo_1, tmem_scores0, 135267472, 1, 0xC0004010U, 0xC0004010U); + elect_commit(c_empty0_addr); + } + elect_commit(score_full0_addr); + } + elect_commit(x_empty0_addr); + } + } + // ---- Role: compute0 ---- + } else if (warp >= 4 && warp <= 7) { + { // compute0_main + int total_tiles_4 = B * num_n_tiles; + int owner_tiles_3 = K_tiles / 2; + int warp_id_in_role = (warp - 4); + int row = warp_id_in_role * 32 + lane; + int row_base = warp_id_in_role * 32 << 16; + unsigned int _phase_score_full0_0 = 0; + unsigned int _phase_slot1_keys_ready_0 = 0; + #pragma unroll 1 + for (unsigned int tile_4 = bid; tile_4 < total_tiles_4; tile_4 += num_bids) { + int batch_2 = tile_4 / (unsigned int)num_n_tiles; + int nt = tile_4 % (unsigned int)num_n_tiles; + float best = -3.4e+38f; + int best_idx = 0; + #pragma unroll 1 + for (int owner_kt_3 = 0; owner_kt_3 < owner_tiles_3; owner_kt_3++) { + int kt_2 = owner_kt_3 * 2; + int kb = kt_2 * 64; + mbarrier_wait(score_full0_addr, _phase_score_full0_0); + _phase_score_full0_0 ^= 1; + float _tmem_load_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x64.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, [%64];" + : "=f"(_tmem_load_0[0]), "=f"(_tmem_load_0[1]), "=f"(_tmem_load_0[2]), "=f"(_tmem_load_0[3]), "=f"(_tmem_load_0[4]), "=f"(_tmem_load_0[5]), "=f"(_tmem_load_0[6]), "=f"(_tmem_load_0[7]), "=f"(_tmem_load_0[8]), "=f"(_tmem_load_0[9]), "=f"(_tmem_load_0[10]), "=f"(_tmem_load_0[11]), "=f"(_tmem_load_0[12]), "=f"(_tmem_load_0[13]), "=f"(_tmem_load_0[14]), "=f"(_tmem_load_0[15]), "=f"(_tmem_load_0[16]), "=f"(_tmem_load_0[17]), "=f"(_tmem_load_0[18]), "=f"(_tmem_load_0[19]), "=f"(_tmem_load_0[20]), "=f"(_tmem_load_0[21]), "=f"(_tmem_load_0[22]), "=f"(_tmem_load_0[23]), "=f"(_tmem_load_0[24]), "=f"(_tmem_load_0[25]), "=f"(_tmem_load_0[26]), "=f"(_tmem_load_0[27]), "=f"(_tmem_load_0[28]), "=f"(_tmem_load_0[29]), "=f"(_tmem_load_0[30]), "=f"(_tmem_load_0[31]), "=f"(_tmem_load_0[32]), "=f"(_tmem_load_0[33]), "=f"(_tmem_load_0[34]), "=f"(_tmem_load_0[35]), "=f"(_tmem_load_0[36]), "=f"(_tmem_load_0[37]), "=f"(_tmem_load_0[38]), "=f"(_tmem_load_0[39]), "=f"(_tmem_load_0[40]), "=f"(_tmem_load_0[41]), "=f"(_tmem_load_0[42]), "=f"(_tmem_load_0[43]), "=f"(_tmem_load_0[44]), "=f"(_tmem_load_0[45]), "=f"(_tmem_load_0[46]), "=f"(_tmem_load_0[47]), "=f"(_tmem_load_0[48]), "=f"(_tmem_load_0[49]), "=f"(_tmem_load_0[50]), "=f"(_tmem_load_0[51]), "=f"(_tmem_load_0[52]), "=f"(_tmem_load_0[53]), "=f"(_tmem_load_0[54]), "=f"(_tmem_load_0[55]), "=f"(_tmem_load_0[56]), "=f"(_tmem_load_0[57]), "=f"(_tmem_load_0[58]), "=f"(_tmem_load_0[59]), "=f"(_tmem_load_0[60]), "=f"(_tmem_load_0[61]), "=f"(_tmem_load_0[62]), "=f"(_tmem_load_0[63]) + : "r"(taddr + (unsigned int)row_base) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 8, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty0_addr); + } + #pragma unroll + for (int kk = 0; kk < 64; kk++) { + float score = _tmem_load_0[kk] - 0.5f * c_sq[batch_2 * K + kb + kk]; + if (score > best) { + best = score; + best_idx = kb + kk; + } + } + } + uint32_t _amf_u_0 = __float_as_uint(best); + uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; + unsigned int _amf_enc_0 = _amf_u_0 ^ _amf_mask_0; + slot0_keys[row] = (unsigned long long)_amf_enc_0 << 32 | 4294967295 - (unsigned long long)best_idx; + asm volatile("barrier.sync 10, %0;" :: "r"(128)); + mbarrier_wait(slot1_keys_ready_addr, _phase_slot1_keys_ready_0); + _phase_slot1_keys_ready_0 ^= 1; + unsigned long long peer_key = slot1_keys[row]; + unsigned long long own_key = slot0_keys[row]; + unsigned long long best_key = ((peer_key > own_key) ? peer_key : own_key); + int idx = (int)(4294967295 - (best_key & 4294967295)); + *((int*)(out + (batch_2 * N + nt * 128 + row))) = idx; + } + } + // ---- Role: compute1 ---- + } else if (warp >= 8 && warp <= 11) { + { // compute1_main + int total_tiles_5 = B * num_n_tiles; + int owner_tiles_4 = K_tiles / 2; + int warp_id_in_role_1 = (warp - 8); + int row_1 = warp_id_in_role_1 * 32 + lane; + int row_base_1 = warp_id_in_role_1 * 32 << 16; + unsigned int _phase_score_full1_0 = 0; + #pragma unroll 1 + for (unsigned int tile_5 = bid; tile_5 < total_tiles_5; tile_5 += num_bids) { + int batch_3 = tile_5 / (unsigned int)num_n_tiles; + float best_1 = -3.4e+38f; + int best_idx_1 = 64; + #pragma unroll 1 + for (int owner_kt_4 = 0; owner_kt_4 < owner_tiles_4; owner_kt_4++) { + int kt_3 = owner_kt_4 * 2 + 1; + int kb_1 = kt_3 * 64; + mbarrier_wait(score_full1_addr, _phase_score_full1_0); + _phase_score_full1_0 ^= 1; + float _tmem_load_1[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x64.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, [%64];" + : "=f"(_tmem_load_1[0]), "=f"(_tmem_load_1[1]), "=f"(_tmem_load_1[2]), "=f"(_tmem_load_1[3]), "=f"(_tmem_load_1[4]), "=f"(_tmem_load_1[5]), "=f"(_tmem_load_1[6]), "=f"(_tmem_load_1[7]), "=f"(_tmem_load_1[8]), "=f"(_tmem_load_1[9]), "=f"(_tmem_load_1[10]), "=f"(_tmem_load_1[11]), "=f"(_tmem_load_1[12]), "=f"(_tmem_load_1[13]), "=f"(_tmem_load_1[14]), "=f"(_tmem_load_1[15]), "=f"(_tmem_load_1[16]), "=f"(_tmem_load_1[17]), "=f"(_tmem_load_1[18]), "=f"(_tmem_load_1[19]), "=f"(_tmem_load_1[20]), "=f"(_tmem_load_1[21]), "=f"(_tmem_load_1[22]), "=f"(_tmem_load_1[23]), "=f"(_tmem_load_1[24]), "=f"(_tmem_load_1[25]), "=f"(_tmem_load_1[26]), "=f"(_tmem_load_1[27]), "=f"(_tmem_load_1[28]), "=f"(_tmem_load_1[29]), "=f"(_tmem_load_1[30]), "=f"(_tmem_load_1[31]), "=f"(_tmem_load_1[32]), "=f"(_tmem_load_1[33]), "=f"(_tmem_load_1[34]), "=f"(_tmem_load_1[35]), "=f"(_tmem_load_1[36]), "=f"(_tmem_load_1[37]), "=f"(_tmem_load_1[38]), "=f"(_tmem_load_1[39]), "=f"(_tmem_load_1[40]), "=f"(_tmem_load_1[41]), "=f"(_tmem_load_1[42]), "=f"(_tmem_load_1[43]), "=f"(_tmem_load_1[44]), "=f"(_tmem_load_1[45]), "=f"(_tmem_load_1[46]), "=f"(_tmem_load_1[47]), "=f"(_tmem_load_1[48]), "=f"(_tmem_load_1[49]), "=f"(_tmem_load_1[50]), "=f"(_tmem_load_1[51]), "=f"(_tmem_load_1[52]), "=f"(_tmem_load_1[53]), "=f"(_tmem_load_1[54]), "=f"(_tmem_load_1[55]), "=f"(_tmem_load_1[56]), "=f"(_tmem_load_1[57]), "=f"(_tmem_load_1[58]), "=f"(_tmem_load_1[59]), "=f"(_tmem_load_1[60]), "=f"(_tmem_load_1[61]), "=f"(_tmem_load_1[62]), "=f"(_tmem_load_1[63]) + : "r"(taddr + (unsigned int)row_base_1 + 64) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + asm volatile("barrier.sync 9, %0;" :: "r"(128)); + if (elect_sync()) { + mbarrier_arrive(score_empty1_addr); + } + #pragma unroll + for (int kk_1 = 0; kk_1 < 64; kk_1++) { + float score_1 = _tmem_load_1[kk_1] - 0.5f * c_sq[batch_3 * K + kb_1 + kk_1]; + if (score_1 > best_1) { + best_1 = score_1; + best_idx_1 = kb_1 + kk_1; + } + } + } + uint32_t _amf_u_0 = __float_as_uint(best_1); + uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; + unsigned int _amf_enc_1 = _amf_u_0 ^ _amf_mask_0; + slot1_keys[row_1] = (unsigned long long)_amf_enc_1 << 32 | 4294967295 - (unsigned long long)best_idx_1; + asm volatile("barrier.sync 11, %0;" :: "r"(128)); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + if (elect_sync()) { + mbarrier_arrive(slot1_keys_ready_addr); + } + } + } + // ---- Role: mma1 ---- + } else if (warp == 12) { + { // mma1_main + int total_tiles_6 = B * num_n_tiles; + int owner_tiles_5 = K_tiles / 2; + unsigned int _phase_x_ready1_0 = 0; + unsigned int _phase_score_empty1_0 = 1; + unsigned int _phase_c_full1_0 = 0; + #pragma unroll 1 + for (unsigned int tile_6 = bid; tile_6 < total_tiles_6; tile_6 += num_bids) { + mbarrier_wait(x_ready1_addr, _phase_x_ready1_0); + _phase_x_ready1_0 ^= 1; + #pragma unroll 1 + for (int owner_kt_5 = 0; owner_kt_5 < owner_tiles_5; owner_kt_5++) { + mbarrier_wait(score_empty1_addr, _phase_score_empty1_0); + _phase_score_empty1_0 ^= 1; + mbarrier_wait(c_full1_addr, _phase_c_full1_0); + _phase_c_full1_0 ^= 1; + asm volatile("fence.proxy.async;"); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_2 = sx_addr; + int _mma_a_lo_2 = make_warp_uniform((_mma_a_addr_2 >> 4) & 0x3FFF); + int _mma_b_addr_2 = sc1_addr; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + mma_ss_step(_mma_a_lo_2, _mma_b_lo_2, tmem_scores1, 135267472, 0, 0xC0004010U, 0xC0004010U); + elect_commit(c_empty1_addr); + #pragma unroll + for (int ft_5 = 1; ft_5 < 7; ft_5++) { + mbarrier_wait(c_full1_addr, _phase_c_full1_0); + _phase_c_full1_0 ^= 1; + asm volatile("fence.proxy.async;"); + asm volatile("tcgen05.fence::after_thread_sync;"); + int _mma_a_addr_3 = sx_addr + (unsigned int)(ft_5 * 4096); + int _mma_a_lo_3 = make_warp_uniform((_mma_a_addr_3 >> 4) & 0x3FFF); + int _mma_b_addr_3 = sc1_addr; + int _mma_b_lo_3 = make_warp_uniform((_mma_b_addr_3 >> 4) & 0x3FFF); + mma_ss_step(_mma_a_lo_3, _mma_b_lo_3, tmem_scores1, 135267472, 1, 0xC0004010U, 0xC0004010U); + elect_commit(c_empty1_addr); + } + elect_commit(score_full1_addr); + } + elect_commit(x_empty1_addr); + } + } + } + + // Cleanup + __syncthreads(); // barrier before TMEM dealloc + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(128)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0356.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0356.cu new file mode 100644 index 00000000..7ed658bb --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0356.cu @@ -0,0 +1,562 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_X_RAW_OFF 1024 +#define SMEM_X_RAW_STAGE_BYTES 14336 +#define SMEM_X_RAW_STRIDE 14336 +#define SMEM_C_DIRECT00_OFF 15360 +#define SMEM_C_DIRECT00_STAGE_BYTES 1792 +#define SMEM_C_DIRECT00_STRIDE 1792 +#define SMEM_C_DIRECT01_OFF 17152 +#define SMEM_C_DIRECT01_STAGE_BYTES 1792 +#define SMEM_C_DIRECT01_STRIDE 1792 +#define SMEM_C_DIRECT10_OFF 18944 +#define SMEM_C_DIRECT10_STAGE_BYTES 1792 +#define SMEM_C_DIRECT10_STRIDE 1792 +#define SMEM_C_DIRECT11_OFF 20736 +#define SMEM_C_DIRECT11_STAGE_BYTES 1792 +#define SMEM_C_DIRECT11_STRIDE 1792 +#define SMEM_SX_OFF 22528 +#define SMEM_SX_STAGE_BYTES 14336 +#define SMEM_SX_STRIDE 14336 +#define SMEM_SS_OFF 36864 +#define SMEM_SS_STAGE_BYTES 4096 +#define SMEM_SS_STRIDE 4096 +#define SMEM_GROUP_KEYS_OFF 40960 +#define SMEM_GROUP_KEYS_STAGE_BYTES 1024 +#define SMEM_GROUP_KEYS_STRIDE 1024 +#define SMEM_LOCAL_KEYS_OFF 41984 +#define SMEM_LOCAL_KEYS_STAGE_BYTES 512 +#define SMEM_LOCAL_KEYS_STRIDE 512 +#define SMEM_CLUSTER_KEYS_OFF 42496 +#define SMEM_CLUSTER_KEYS_STAGE_BYTES 4096 +#define SMEM_CLUSTER_KEYS_STRIDE 4096 +#define SMEM_TOTAL 46592 +#define THREADS 256 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ uint32_t smem_addr(const void* ptr) { + uint32_t addr; + asm("{\n\t" + ".reg .u64 u64addr;\n\t" + "cvta.to.shared.u64 u64addr, %1;\n\t" + "cvt.u32.u64 %0, u64addr;\n\t" + "}\n" : "=r"(addr) : "l"(ptr)); + return addr; +} + + +__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) { + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(remote) : "r"(local_addr), "r"(rank)); + return remote; +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ void cp_async_bulk_gmem2smem( + unsigned smem_addr, const void* gmem_ptr, unsigned bytes, int mbar_addr) { + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(smem_addr), "l"(gmem_ptr), "r"(bytes), "r"(mbar_addr) + : "memory"); +} + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, float* __restrict__ c_sq, int* __restrict__ out, int B, int N, int D, int K, int num_n_tiles) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 8; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 8; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* x_raw = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int x_raw_addr = smem + 1024; + __nv_bfloat16* c_direct00 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 15360); + const int c_direct00_addr = smem + 15360; + __nv_bfloat16* c_direct01 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17152); + const int c_direct01_addr = smem + 17152; + __nv_bfloat16* c_direct10 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 18944); + const int c_direct10_addr = smem + 18944; + __nv_bfloat16* c_direct11 = reinterpret_cast<__nv_bfloat16*>(smem_raw + 20736); + const int c_direct11_addr = smem + 20736; + __nv_bfloat16* sx = reinterpret_cast<__nv_bfloat16*>(smem_raw + 22528); + const int sx_addr = smem + 22528; + float* ss = reinterpret_cast(smem_raw + 36864); + const int ss_addr = smem + 36864; + unsigned long long* group_keys = reinterpret_cast(smem_raw + 40960); + const int group_keys_addr = smem + 40960; + unsigned long long* local_keys = reinterpret_cast(smem_raw + 41984); + const int local_keys_addr = smem + 41984; + unsigned long long* cluster_keys = reinterpret_cast(smem_raw + 42496); + const int cluster_keys_addr = smem + 42496; + + // Mbarrier init (6 groups, 13 barriers) + // Mbarriers at smem_raw[0..104) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // x_ready: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + // c_ready00: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 8, 1, leader); + // c_ready01: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 16, 1, leader); + // c_ready10: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 24, 1, leader); + // c_ready11: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 32, 1, leader); + // keys_ready: 8 barriers, init_count=8 + mbarrier_init_pred(smem + 40, 8, leader); + mbarrier_init_pred(smem + 48, 8, leader); + mbarrier_init_pred(smem + 56, 8, leader); + mbarrier_init_pred(smem + 64, 8, leader); + mbarrier_init_pred(smem + 72, 8, leader); + mbarrier_init_pred(smem + 80, 8, leader); + mbarrier_init_pred(smem + 88, 8, leader); + mbarrier_init_pred(smem + 96, 8, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); + + const int mbar_base = smem; + #define x_ready_addr (mbar_base + 0) + #define c_ready00_addr (mbar_base + 8) + #define c_ready01_addr (mbar_base + 16) + #define c_ready10_addr (mbar_base + 24) + #define c_ready11_addr (mbar_base + 32) + #define keys_ready_addr (mbar_base + 40) + + // === Task calls (dependency order) === + int total_tiles = B * num_n_tiles; + unsigned int _phase_x_ready_0 = 0; + unsigned int _phase_c_ready00_0 = 0; + unsigned int _phase_c_ready01_0 = 0; + unsigned int _phase_c_ready10_0 = 0; + unsigned int _phase_c_ready11_0 = 0; + unsigned int _phase_keys_ready_0 = 0; + #pragma unroll 1 + for (unsigned int tile = cluster_id; tile < total_tiles; tile += num_clusters) { + int batch = tile / (unsigned int)num_n_tiles; + int nt = tile % (unsigned int)num_n_tiles; + int point_base = batch * N + nt * 64; + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(x_ready_addr, 14336); + cp_async_bulk_gmem2smem(x_raw_addr, reinterpret_cast(reinterpret_cast(x) + ((unsigned long long)(point_base * D) * (unsigned long long)2)), 14336, x_ready_addr); + } + } + mbarrier_wait(x_ready_addr, _phase_x_ready_0); + _phase_x_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + #pragma unroll 1 + for (unsigned int i = tid; i < 7168; i += 256) { + int row = i / 112; + int col = i % 112; + { + __nv_bfloat16 _bval_3404531696 = __float2bfloat16_rn(x_raw[i]); + uint16_t _bits_3404531696 = *(uint16_t*)&_bval_3404531696; + uint32_t _addr_3404531696 = static_cast((sx_addr + (unsigned int)(row * 224 + col * 2))); + asm volatile("st.shared.b16 [%0], %1;" :: "r"(_addr_3404531696), "h"(_bits_3404531696) : "memory"); + } + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + int warp_id_in_role = (warp - 0); + int group = warp_id_in_role / 4; + int local_warp = warp_id_in_role % 4; + int local_row = lane % 16; + int row_base = local_warp * 16; + float best = -3.4e+38f; + int owner_k_tiles = K / 64; + int group_k_tiles = owner_k_tiles / 2; + int group_tile_begin = cta_rank * owner_k_tiles + group * group_k_tiles; + int best_idx = group_tile_begin * 8; + int first_kbase = group_tile_begin * 8; + if (group == 0) { + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready00_addr, 1792); + cp_async_bulk_gmem2smem(c_direct00_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + first_kbase) * D) * (unsigned long long)2)), 1792, c_ready00_addr); + } + } + } else if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready10_addr, 1792); + cp_async_bulk_gmem2smem(c_direct10_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + first_kbase) * D) * (unsigned long long)2)), 1792, c_ready10_addr); + } + } + #pragma unroll 1 + for (int group_kt = 0; group_kt < group_k_tiles; group_kt++) { + int current_stage = group_kt % 2; + int kt = group_tile_begin + group_kt; + int kbase = kt * 8; + bool has_next = group_k_tiles > group_kt + 1; + if (has_next) { + int next_stage = (group_kt + 1) % 2; + int next_kbase = (kt + 1) * 8; + if (group == 0) { + if (next_stage == 0) { + if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready00_addr, 1792); + cp_async_bulk_gmem2smem(c_direct00_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready00_addr); + } + } + } else if (warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready01_addr, 1792); + cp_async_bulk_gmem2smem(c_direct01_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready01_addr); + } + } + } else if (next_stage == 0) { + if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready10_addr, 1792); + cp_async_bulk_gmem2smem(c_direct10_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready10_addr); + } + } + } else { + if (warp == 4) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(c_ready11_addr, 1792); + cp_async_bulk_gmem2smem(c_direct11_addr, reinterpret_cast(reinterpret_cast(centroids) + ((unsigned long long)((batch * K + next_kbase) * D) * (unsigned long long)2)), 1792, c_ready11_addr); + } + } + } + } + if (group == 0) { + if (current_stage == 0) { + mbarrier_wait(c_ready00_addr, _phase_c_ready00_0); + _phase_c_ready00_0 ^= 1; + } else { + mbarrier_wait(c_ready01_addr, _phase_c_ready01_0); + _phase_c_ready01_0 ^= 1; + } + } else if (current_stage == 0) { + mbarrier_wait(c_ready10_addr, _phase_c_ready10_0); + _phase_c_ready10_0 ^= 1; + } else { + mbarrier_wait(c_ready11_addr, _phase_c_ready11_0); + _phase_c_ready11_0 ^= 1; + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + unsigned int a[4]; + unsigned int b[2]; + float acc[4]; + unsigned int a_addr = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + lane / 16 * 16)); + unsigned int b_addr = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (lane / 8 & 1) * 8 * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[0]), "=f"(acc[1]), "=f"(acc[2]), "=f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_0 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 32))); + unsigned int b_addr_1 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (16 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_0) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_1) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_2 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 64))); + unsigned int b_addr_3 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (32 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_2) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_3) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_4 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 96))); + unsigned int b_addr_5 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (48 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_4) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_5) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_6 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 128))); + unsigned int b_addr_7 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (64 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_6) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_7) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_8 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 160))); + unsigned int b_addr_9 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (80 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_8) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_9) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + unsigned int a_addr_10 = (sx_addr + (unsigned int)((row_base + lane % 16) * 224 + (lane / 16 * 16 + 192))); + unsigned int b_addr_11 = ((group == 0) ? ((current_stage == 0) ? (c_direct00_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)) : (c_direct01_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2))) : ((current_stage == 0) ? (c_direct10_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)) : (c_direct11_addr + (unsigned int)(lane % 8 * 224 + (96 + (lane / 8 & 1) * 8) * 2)))); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) + : "r"(a_addr_10) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(b[0]), "=r"(b[1]) + : "r"(b_addr_11) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); + #pragma unroll + for (int rp = 0; rp < 2; rp++) { + #pragma unroll + for (int cp = 0; cp < 2; cp++) { + int rr = lane / 4 + rp * 8; + int cc = lane % 4 * 2 + cp; + { + uint32_t _addr_3405188304 = static_cast((ss_addr + (unsigned int)((group * 64 + row_base + rr) * 32 + cc * 4))); + asm volatile("st.shared.f32 [%0], %1;" :: "r"(_addr_3405188304), "f"(acc[rp * 2 + cp]) : "memory"); + } + } + } + __syncwarp(); + if (lane < 16) { + #pragma unroll + for (int kk = 0; kk < 8; kk++) { + float score = ss[(group * 64 + row_base + local_row) * 8 + kk] - 0.5f * c_sq[batch * K + kbase + kk]; + if (score > best) { + best = score; + best_idx = kbase + kk; + } + } + } + if (group == 0) { + asm volatile("barrier.sync 1, %0;" :: "r"(128)); + } else { + asm volatile("barrier.sync 2, %0;" :: "r"(128)); + } + } + if (lane < 16) { + uint32_t _amf_u_0 = __float_as_uint(best); + uint32_t _amf_mask_0 = -int32_t(_amf_u_0 >> 31) | 0x80000000u; + unsigned int _amf_enc_0 = _amf_u_0 ^ _amf_mask_0; + unsigned long long shift64 = 32; + unsigned long long mask64 = 4294967295; + group_keys[group * 64 + row_base + local_row] = (unsigned long long)_amf_enc_0 << shift64 | mask64 - (unsigned long long)best_idx; + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + if (group == 0 && lane < 16) { + unsigned long long first_key = group_keys[row_base + local_row]; + unsigned long long second_key = group_keys[64 + row_base + local_row]; + local_keys[row_base + local_row] = ((second_key > first_key) ? second_key : first_key); + } + asm volatile("barrier.sync 3, %0;" :: "r"(256)); + if (warp == 0) { + if (elect_sync()) { + uint32_t _mapa_0; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_0) : "r"(cluster_keys_addr + (unsigned int)(cta_rank * 512)), "r"(0)); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + uint32_t _mapa_1; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_1) : "r"(keys_ready_addr), "r"(0)); + asm volatile( + "mbarrier.arrive.expect_tx.release.cluster.shared::cluster.b64 _, [%0], %1;" + :: "r"(_mapa_1), "r"((uint32_t)(512)) : "memory"); + uint32_t _mapa_2; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mapa_2) : "r"(keys_ready_addr), "r"(0)); + asm volatile( + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(_mapa_0), "r"(local_keys_addr), "r"((uint32_t)(512)), "r"(_mapa_2) + : "memory"); + } + } + if (cta_rank == 0) { + mbarrier_wait(keys_ready_addr, _phase_keys_ready_0); + _phase_keys_ready_0 ^= 1; + asm volatile("tcgen05.fence::after_thread_sync;"); + int row_1 = tid / 2; + int lane_pair = tid % 2; + if (row_1 < 64) { + unsigned long long best_key = 0; + #pragma unroll + for (int peer = lane_pair; peer < 8; peer += 2) { + unsigned long long key = cluster_keys[peer * 64 + row_1]; + if (key > best_key) { + best_key = key; + } + } + unsigned long long _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, best_key, 1); + unsigned long long peer_key = _shfl_xor_0; + if (peer_key > best_key) { + best_key = peer_key; + } + if (lane_pair == 0) { + unsigned long long mask64_1 = 4294967295; + int idx = (int)(mask64_1 - (best_key & mask64_1)); + *((int*)(out + (batch * N + nt * 64 + row_1))) = idx; + } + } + } + } + + // Cleanup + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0357.cu b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0357.cu new file mode 100644 index 00000000..ea986bdf --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/cuda/dispatch_kernel_0357.cu @@ -0,0 +1,138 @@ +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" + : "=r"(result) : "r"(x)); + return result; +} + +#define LOOM_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define THREADS 256 + +#include + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1(__nv_bfloat16* __restrict__ x, __nv_bfloat16* __restrict__ centroids, __nv_bfloat16* __restrict__ x_pad, __nv_bfloat16* __restrict__ c_pad, int B, int N, int D, int K, int total_x_pad, int total_c_pad) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // === Task calls (dependency order) === + int grid_stride = num_bids * 256; + int start = bid * 256 + tid; + int total_x_vecs = total_x_pad / 8; + int total_c_vecs = total_c_pad / 8; + #pragma unroll 1 + for (unsigned int vec_idx = start; vec_idx < total_x_vecs; vec_idx += grid_stride) { + int d_pad = vec_idx % 24 * 8; + int row = vec_idx / 24; + int dst_base = vec_idx * 8; + if (d_pad < D) { + float _vec_load_0[8]; + { + const uint4* _vptr_0 = reinterpret_cast(x + (row * D + d_pad) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + __nv_bfloat16* _velems_0 = reinterpret_cast<__nv_bfloat16*>(&_vld_0[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_0[0 + _blk * 8 + _j] = __bfloat162float(_velems_0[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_0[0 + 0], _vec_load_0[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_0[0 + 2], _vec_load_0[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_0[0 + 4], _vec_load_0[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_0[0 + 6], _vec_load_0[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros[8]; + zeros[0] = 0.0f; + zeros[1] = 0.0f; + zeros[2] = 0.0f; + zeros[3] = 0.0f; + zeros[4] = 0.0f; + zeros[5] = 0.0f; + zeros[6] = 0.0f; + zeros[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(x_pad + dst_base))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + #pragma unroll 1 + for (unsigned int vec_idx_1 = start; vec_idx_1 < total_c_vecs; vec_idx_1 += grid_stride) { + int d_pad_1 = vec_idx_1 % 24 * 8; + int row_1 = vec_idx_1 / 24; + int dst_base_1 = vec_idx_1 * 8; + if (d_pad_1 < D) { + float _vec_load_1[8]; + { + const uint4* _vptr_1 = reinterpret_cast(centroids + (row_1 * D + d_pad_1) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + __nv_bfloat16* _velems_1 = reinterpret_cast<__nv_bfloat16*>(&_vld_1[_blk]); + #pragma unroll + for (int _j = 0; _j < 8; _j++) + _vec_load_1[0 + _blk * 8 + _j] = __bfloat162float(_velems_1[_j]); + } + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(_vec_load_1[0 + 0], _vec_load_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_vec_load_1[0 + 2], _vec_load_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_vec_load_1[0 + 4], _vec_load_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_vec_load_1[0 + 6], _vec_load_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + float zeros_1[8]; + zeros_1[0] = 0.0f; + zeros_1[1] = 0.0f; + zeros_1[2] = 0.0f; + zeros_1[3] = 0.0f; + zeros_1[4] = 0.0f; + zeros_1[5] = 0.0f; + zeros_1[6] = 0.0f; + zeros_1[7] = 0.0f; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros_1[0 + 0], zeros_1[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros_1[0 + 2], zeros_1[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros_1[0 + 4], zeros_1[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros_1[0 + 6], zeros_1[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(c_pad + dst_base_1))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } +} + +} // extern "C" + diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/interface.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/interface.py new file mode 100644 index 00000000..87fed6c0 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/interface.py @@ -0,0 +1,1085 @@ +from __future__ import annotations + +import math +from collections import OrderedDict +from contextlib import contextmanager +from dataclasses import dataclass, field +from functools import wraps +from threading import Condition, Event, RLock +from typing import Any + +from ._dispatch import flash_kmeans_assign_dispatcher as _dispatcher +from ._dispatch_runtime import dispatch_launch_options +from ._launch_plan import GraphCaptureUnsupported, build_graph_exec_plan, build_launch_plan +from ._runtime import detect_gpu_arch, runtime_activity_snapshot + +SEMANTIC_ENTRYPOINT = "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval" +_PREPARE_LOCK = RLock() + + +@dataclass(frozen=True) +class PreparedFlashKMeansAssign: + """Exact tensors, workspace, output, and direct launch plan for KMeans.""" + + inputs: dict[str, Any] + launch_plan: Any + stream: Any = None + timeout_ms: float | None = None + + @property + def out(self) -> Any: + return self.inputs["out"] + + @property + def selected_route(self) -> str: + return self.launch_plan.selected_route + + +@dataclass(frozen=True) +class _ResolvedRoute: + """Frozen-cascade decision adapted to the shared plan-constructor contract. + + ``launch`` submits through the root dispatcher, which re-derives exactly + this decision from the same frozen guard cascade — the cascade remains the + single source of routing truth and runs once per signature, at plan + construction (or per call for host-data-dependent routes). + """ + + route_id: str + launch_entrypoint: str + exact_contract: bool = True + + def launch( + self, + inputs: dict[str, Any], + *, + stream: Any = None, + timeout_ms: float | None = None, + ) -> Any: + with dispatch_launch_options(stream=stream, timeout_ms=timeout_ms): + return _dispatcher.launch_for_eval(inputs) + + +@dataclass +class _RuntimeSlot: + """One stream-bound, signature-specialized LaunchPlan plus support state. + + ``plan`` freezes the resolved route's leaf launches — including any + zero-pad staging kernels the route delegates through — or keeps the cached + per-call launcher for host-data-dependent routes. ``norm_plan`` is the + route-required fused BF16 pair row-norm support launch, bound by pointer + overwrite before the plan; its outputs are the slot-owned ``x_sq``/``c_sq`` + scratch with stable pointers. ``graph`` is the signature's captured CUDA + graph over the full norm + route kernel chain; when set, a hot call binds + pointers host-side and replays the graph instead of submitting the + prepared launches one by one (``graph_capture_error`` records why a frozen + plan stayed on the prepared path). ``default_outputs`` alternates between + two plan-owned tensors so consecutive default-output calls never hand the + caller the allocation the previous call returned. ``default_flip`` is only + read or written while ``lock`` is held. + """ + + plan: Any + route: Any + norm_plan: Any + x_sq: Any + c_sq: Any + internal_x_sq: bool + internal_c_sq: bool + norm_compute_fields: tuple[str, ...] + default_outputs: tuple[Any, Any] + graph: Any = None + graph_capture_error: str | None = None + lock: RLock = field(default_factory=RLock, repr=False) + default_flip: int = 0 + + +@dataclass +class _PendingPreparation: + event: Event = field(default_factory=Event, repr=False) + error: BaseException | None = field(default=None, repr=False) + + +def _guard_runtime_compute(method: Any) -> Any: + """Keep clear() from crossing an in-progress lookup/rebind/enqueue call.""" + + @wraps(method) + def guarded(self: Any, *args: Any, **kwargs: Any) -> Any: + with self._lifecycle: + while self._clearing: + self._lifecycle.wait() + self._active_computes += 1 + try: + return method(self, *args, **kwargs) + finally: + with self._lifecycle: + self._active_computes -= 1 + if self._active_computes == 0: + self._lifecycle.notify_all() + + return guarded + + +def _normalize_device_index(torch: Any, device: Any) -> int: + if device is None: + return int(torch.cuda.current_device()) + if isinstance(device, bool): + raise TypeError("device must identify a CUDA device") + if isinstance(device, int): + if device < 0: + raise ValueError("CUDA device index must be non-negative") + return int(device) + device_type = getattr(device, "type", None) + device_index = getattr(device, "index", None) + if device_type is None: + device_ctor = getattr(torch, "device", None) + if device_ctor is None: + raise TypeError("device must identify a CUDA device") + normalized = device_ctor(device) + device_type = getattr(normalized, "type", None) + device_index = getattr(normalized, "index", None) + if device_type != "cuda": + raise ValueError(f"Flash-KMeans runtime requires a CUDA device, got {device_type!r}") + resolved_index = int(torch.cuda.current_device() if device_index is None else device_index) + if resolved_index < 0: + raise ValueError("CUDA device index must be non-negative") + return resolved_index + + +def _validate_timeout(timeout_ms: float | None) -> float | None: + if timeout_ms is None: + return None + if isinstance(timeout_ms, bool) or not isinstance(timeout_ms, int | float): + raise TypeError("timeout_ms must be a positive finite number or None") + value = float(timeout_ms) + if not math.isfinite(value) or value <= 0: + raise ValueError("timeout_ms must be a positive finite number") + return value + + +def _resolve_stream(torch: Any, device_index: int, stream: Any) -> Any: + resolved_stream = torch.cuda.current_stream(device_index) if stream is None else stream + stream_device = getattr(resolved_stream, "device", None) + stream_device_index = getattr(stream_device, "index", stream_device) + if stream_device_index is not None and int(stream_device_index) != int(device_index): + raise ValueError( + f"Flash-KMeans stream device {stream_device_index} does not match input device {device_index}" + ) + return resolved_stream + + +def _same_cuda_stream(left: Any, right: Any) -> bool: + return left is right or int(left.cuda_stream) == int(right.cuda_stream) + + +@contextmanager +def _device_stream_context( + torch: Any, + device_index: int, + stream: Any, + *, + context_is_active: bool, +): + """Resolve one stream and enter it only when the caller has not already done so.""" + + if context_is_active: + if stream is None: + raise RuntimeError("an active Flash-KMeans CUDA context requires an explicit stream") + yield stream + return + with torch.cuda.device(device_index): + resolved_stream = _resolve_stream(torch, device_index, stream) + with torch.cuda.stream(resolved_stream): + yield resolved_stream + + +def _signature_alias_topology( + x: Any, + centroids: Any, + x_sq: Any, + c_sq: Any, + out: Any, +) -> tuple[int, ...]: + """Canonical pointer-equivalence classes independent of concrete addresses. + + A default output comes from the signature slot's plan-owned pool, which + can never alias a live caller tensor, so ``out is None`` contributes a + fresh class without reading any pointer. + """ + + classes: dict[int, int] = {} + topology: list[int] = [] + for tensor in (x, centroids, x_sq, c_sq): + if tensor is None: + topology.append(-1) + continue + pointer = int(tensor.data_ptr()) + if pointer not in classes: + classes[pointer] = len(classes) + topology.append(classes[pointer]) + if out is None: + topology.append(len(classes)) + else: + pointer = int(out.data_ptr()) + if pointer not in classes: + classes[pointer] = len(classes) + topology.append(classes[pointer]) + return tuple(topology) + + +def _require_owned_output(result: Any, inputs: dict[str, Any]) -> None: + """Reject routes whose hot path would require an uncaptured output copy.""" + + produced = result.get("cluster_ids") if isinstance(result, dict) else result + if produced is None: + return + out = inputs["out"] + if produced is not out and int(produced.data_ptr()) != int(out.data_ptr()): + raise RuntimeError( + "prepared Flash-KMeans route must write the caller-owned output through captured launches" + ) + + +def _validate_compute_tensors( + torch: Any, + x: Any, + centroids: Any, + *, + device_index: int, +) -> tuple[int, int, int, int]: + """Validate the public compute ABI without allocating anything.""" + + if not all(isinstance(item, torch.Tensor) and item.is_cuda for item in (x, centroids)): + raise TypeError("x and centroids must be CUDA torch.Tensor objects") + if x.dtype is not torch.bfloat16 or centroids.dtype is not torch.bfloat16: + raise TypeError("x and centroids dtype must be bfloat16") + if x.ndim != 3 or centroids.ndim != 3 or not x.is_contiguous() or not centroids.is_contiguous(): + raise ValueError("x and centroids must be contiguous [B, rows, D] tensors") + bsz, n_points, dim = map(int, x.shape) + c_bsz, n_clusters, c_dim = map(int, centroids.shape) + if (bsz, dim) != (c_bsz, c_dim) or x.device != centroids.device: + raise ValueError("x and centroids batch/feature dimensions and device must match") + input_device_index = x.device.index + if input_device_index is None: + input_device_index = torch.cuda.current_device() + if int(input_device_index) != int(device_index): + raise ValueError( + f"Flash-KMeans runtime targets CUDA device {device_index}, got input device {input_device_index}" + ) + return bsz, n_points, dim, n_clusters + + +class FlashKMeansAssignRuntime: + """Long-lived KMeans assignment runtime with per-shape, per-stream plans. + + ``compute`` accepts new tensors and new shapes. The first call for one + signature runs the frozen guard cascade once, freezes the resolved + route's exact leaf launches (plus the route-required fused pair row-norm + support launch) into a per-signature plan, and captures the full kernel + chain into one CUDA graph. Cache hits overwrite the recorded pointer + carriers in place and replay the graph on the plan's stream — no dispatch + re-evaluation, no argument re-marshalling, no per-launch stream query, + and no per-call default-output allocation (defaults ping-pong between two + plan-owned tensors, so consecutive default-output calls never alias). + Workspace is never shared across CUDA streams. + """ + + def __init__( + self, + *, + device: Any = None, + arch: str | None = None, + timeout_ms: float | None = None, + max_cached_shapes: int | None = None, + compile: str = "lazy", + ) -> None: + import torch + + if compile != "lazy": + raise ValueError("compile must be 'lazy'; use warmup() to eagerly prepare known shapes") + if max_cached_shapes is not None: + if isinstance(max_cached_shapes, bool) or not isinstance(max_cached_shapes, int): + raise TypeError("max_cached_shapes must be a positive integer or None") + if max_cached_shapes <= 0: + raise ValueError("max_cached_shapes must be positive") + self.device_index = _normalize_device_index(torch, device) + with torch.cuda.device(self.device_index): + detected_arch = str(detect_gpu_arch()) + self.arch = detected_arch if arch is None else str(arch) + if self.arch != detected_arch: + raise ValueError( + f"Flash-KMeans runtime arch must match its device: requested {self.arch}, detected {detected_arch}" + ) + self.timeout_ms = _validate_timeout(timeout_ms) + self.max_cached_shapes = max_cached_shapes + self._torch = torch + self._slots: OrderedDict[tuple[Any, ...], _RuntimeSlot] = OrderedDict() + self._preparing: dict[tuple[Any, ...], _PendingPreparation] = {} + self._cache_lock = RLock() + self._lifecycle = Condition(RLock()) + self._active_computes = 0 + self._clearing = False + self._hits = 0 + self._misses = 0 + + def cache_info(self) -> dict[str, int | None]: + with self._cache_lock: + return { + "size": len(self._slots), + "hits": self._hits, + "misses": self._misses, + "max_cached_shapes": self.max_cached_shapes, + } + + @_guard_runtime_compute + def compute( + self, + x: Any, + centroids: Any, + *, + out: Any | None = None, + x_sq: Any | None = None, + c_sq: Any | None = None, + stream: Any = None, + timeout_ms: float | None = None, + return_info: bool = False, + ): + """Run one KMeans assignment through this signature's launch plan. + + A cache miss runs the frozen guard cascade once and freezes its exact + leaf launches into a per-signature plan with a captured CUDA graph. + Cache hits are host-side pointer binding plus one graph replay on the + plan's stream. A caller-provided ``out=`` is written through; default + outputs come from the signature's plan-owned pool, so a default-output + result must be consumed before two further default-output calls of the + same signature reuse its storage. + """ + + torch = self._torch + bsz, n_points, dim, n_clusters = _validate_compute_tensors( + torch, + x, + centroids, + device_index=self.device_index, + ) + effective_timeout_ms = self.timeout_ms if timeout_ms is None else _validate_timeout(timeout_ms) + resolved_stream = _resolve_stream(torch, self.device_index, stream) + stream_handle = int(resolved_stream.cuda_stream) + if out is not None: + _require_aux_tensor( + out, + name="out", + shape=(bsz, n_points), + dtype=torch.int32, + device=x.device, + ) + if x_sq is not None: + _require_aux_tensor( + x_sq, + name="x_sq", + shape=(bsz, n_points), + dtype=torch.float32, + device=x.device, + ) + if c_sq is not None: + _require_aux_tensor( + c_sq, + name="c_sq", + shape=(bsz, n_clusters), + dtype=torch.float32, + device=x.device, + ) + key = ( + self.device_index, + self.arch, + bsz, + n_points, + dim, + n_clusters, + "bfloat16", + x_sq is None, + c_sq is None, + _signature_alias_topology(x, centroids, x_sq, c_sq, out), + stream_handle, + ) + activity_before = runtime_activity_snapshot() if return_info else None + with self._cache_lock: + slot = self._slots.get(key) + if slot is not None: + self._slots.move_to_end(key) + self._hits += 1 + cache_hit = slot is not None + owns_slot_lock = False + if slot is None: + slot, owns_slot_lock = self._create_slot( + key, + x=x, + centroids=centroids, + out=out, + x_sq=x_sq, + c_sq=c_sq, + shape=(bsz, n_points, dim, n_clusters), + stream=resolved_stream, + ) + cache_hit = not owns_slot_lock + activity_after = runtime_activity_snapshot() if return_info else None + if not owns_slot_lock: + slot.lock.acquire() + try: + if slot.plan.stream_handle != stream_handle: + raise RuntimeError("cached Flash-KMeans slot changed its prepared stream") + if out is None: + bound_out = slot.default_outputs[slot.default_flip] + slot.default_flip ^= 1 + else: + bound_out = out + bindings = { + "x": x, + "centroids": centroids, + "x_sq": slot.x_sq if slot.internal_x_sq else x_sq, + "c_sq": slot.c_sq if slot.internal_c_sq else c_sq, + "out": bound_out, + } + try: + # Bind first (tensor-map refresh + pointer overwrite are host + # work), then enqueue the norm support launch, then submit + # the plan's kernels. Binding after the norm would turn a + # fresh-pointer tensor-map re-encode into a GPU inter-kernel + # gap between the norm kernel and the route's first kernel. + # A captured signature does every bind host-side, then one + # graph replay covers the whole norm + route kernel chain. + if slot.graph is not None: + slot.plan.bind_hot(bindings) + if slot.norm_plan is not None: + slot.norm_plan.bind_hot(x, centroids) + slot.graph.submit_hot(timeout_ms=effective_timeout_ms) + else: + slot.plan.bind_hot(bindings) + if slot.norm_plan is not None: + slot.norm_plan.launch_hot(x, centroids) + slot.plan.submit_hot(timeout_ms=effective_timeout_ms) + finally: + # Slot-owned norm scratch and pooled default outputs are + # allocated on the plan's stream and only released through + # clear(); caller-provided tensors may live on another + # stream, so record those for allocator safety even on + # partial submission. + _record_stream(resolved_stream, x, centroids, x_sq, c_sq, out) + finally: + slot.lock.release() + + if not return_info: + return bound_out + norm_mode = ( + "fused_bf16_pair_row_norm:" + ",".join(slot.norm_compute_fields) + if slot.norm_compute_fields + else "route_elided_internal_norms" + if slot.internal_x_sq or slot.internal_c_sq + else "explicit_precomputed" + ) + cold_activity = { + "source_read_occurred": bool(activity_after["source_reads"] > activity_before["source_reads"]), + "nvrtc_compile_occurred": bool( + activity_after["nvrtc_compiles"] > activity_before["nvrtc_compiles"] + ), + "module_load_occurred": bool(activity_after["module_loads"] > activity_before["module_loads"]), + "scratch_allocation_occurred": not cache_hit, + } + return bound_out, { + "semantic_entrypoint": SEMANTIC_ENTRYPOINT, + "selected_route": slot.route.route_id, + "launch_entrypoint": slot.route.launch_entrypoint, + "exact_launch_plan": slot.route.exact_contract, + "runtime_cache_hit": cache_hit, + "assignment_launch_count": int(slot.plan.launch_count), + "norm_launch_count": int(slot.norm_plan is not None), + "norm_compute_fields": list(slot.norm_compute_fields), + "runtime_launch_count": int(slot.plan.launch_count) + int(slot.norm_plan is not None), + "norm_mode": norm_mode, + "hot_launch_path": "cuda_graph" if slot.graph is not None else "prepared_launches", + "graph_kernel_count": None if slot.graph is None else int(slot.graph.launch_count), + "graph_capture_error": slot.graph_capture_error, + "arch": self.arch, + "device_index": self.device_index, + "stream_handle": stream_handle, + "cold_first_call_activity": cold_activity, + } + + def _create_slot( + self, + key: tuple[Any, ...], + *, + x: Any, + centroids: Any, + out: Any | None, + x_sq: Any | None, + c_sq: Any | None, + shape: tuple[int, int, int, int], + stream: Any, + ) -> tuple[_RuntimeSlot, bool]: + """Construct and publish this signature's LaunchPlan (the slow path). + + Returns ``(slot, owns_slot_lock)``. A newly published slot is returned + with its lock held so the constructing call launches first; when + another thread published the slot while this one waited, the existing + slot is returned unlocked and counted as a cache hit. Preparation + never holds the runtime cache lock, so resident hot shapes stay + unblocked while a cold signature captures. + """ + + while True: + with self._cache_lock: + slot = self._slots.get(key) + if slot is not None: + self._slots.move_to_end(key) + self._hits += 1 + return slot, False + pending = self._preparing.get(key) + if pending is None: + if ( + self.max_cached_shapes is not None + and len(self._slots) + len(self._preparing) >= self.max_cached_shapes + ): + raise RuntimeError( + "FlashKMeansAssignRuntime cache is full; call clear() only after in-flight work completes" + ) + pending = _PendingPreparation() + self._preparing[key] = pending + break + pending.event.wait() + if pending.error is not None: + raise RuntimeError("KMeans slot preparation failed in another thread") from pending.error + + slot: _RuntimeSlot | None = None + try: + import torch + + from ._row_norm import prepare_bf16_pair_row_norm + + bsz, n_points, dim, n_clusters = shape + internal_x_sq = x_sq is None + internal_c_sq = c_sq is None + with torch.cuda.device(self.device_index), torch.cuda.stream(stream): + default_outputs = tuple( + torch.empty((bsz, n_points), dtype=torch.int32, device=x.device) for _pair in range(2) + ) + slot_x_sq = ( + torch.empty((bsz, n_points), dtype=torch.float32, device=x.device) + if internal_x_sq + else x_sq + ) + slot_c_sq = ( + torch.empty((bsz, n_clusters), dtype=torch.float32, device=x.device) + if internal_c_sq + else c_sq + ) + capture_out = out if out is not None else default_outputs[0] + inputs = { + "B": bsz, + "N": n_points, + "D": dim, + "K": n_clusters, + "dtype": "bfloat16", + "x": x, + "centroids": centroids, + "x_sq": slot_x_sq, + "c_sq": slot_c_sq, + "out": capture_out, + } + with _PREPARE_LOCK: + decision = _dispatcher.select_route(inputs) + route = _ResolvedRoute( + route_id=str(decision.route_id), + launch_entrypoint=str(decision.entrypoint), + ) + plan = build_launch_plan( + inputs, + stream=stream, + arch=self.arch, + validate_result=_require_owned_output, + route=route, + ) + bound_keys = getattr(plan, "bound_input_keys", None) + bound = None if bound_keys is None else set(bound_keys) + # A frozen plan reports exactly which public inputs its + # captured launches bind; routes whose kernels never read a + # norm skip computing it. A per-call plan re-executes the + # route's host program, so provide every internal norm it + # could consume. + compute_x_sq = internal_x_sq and (bound is None or "x_sq" in bound) + compute_c_sq = internal_c_sq and (bound is None or "c_sq" in bound) + norm_plan = None + if compute_x_sq or compute_c_sq: + norm_plan = prepare_bf16_pair_row_norm( + x, + centroids, + slot_x_sq, + slot_c_sq, + compute_x=compute_x_sq, + compute_c=compute_c_sq, + arch=self.arch, + stream=stream, + ) + # Capture the signature's stable kernel chain (the fused pair + # norm first, then the frozen route launches — the exact hot + # submission order) into one CUDA graph. Per-call routes and + # launch modes without a validated capture path stay on the + # prepared-launch hot path, with the reason recorded; any + # other capture failure is an error, not a fallback. + graph_plan = None + graph_capture_error: str | None = None + try: + graph_plan = build_graph_exec_plan( + plan, + support_launches=() if norm_plan is None else (norm_plan.launch_plan,), + ) + except GraphCaptureUnsupported as unsupported: + graph_capture_error = str(unsupported) + norm_compute_fields = tuple( + name for name, enabled in (("x_sq", compute_x_sq), ("c_sq", compute_c_sq)) if enabled + ) + slot = _RuntimeSlot( + plan=plan, + route=route, + norm_plan=norm_plan, + x_sq=slot_x_sq, + c_sq=slot_c_sq, + internal_x_sq=internal_x_sq, + internal_c_sq=internal_c_sq, + norm_compute_fields=norm_compute_fields, + default_outputs=default_outputs, + graph=graph_plan, + graph_capture_error=graph_capture_error, + ) + slot.lock.acquire() + except BaseException as error: + if slot is not None: + try: + slot.lock.release() + except RuntimeError: + pass + with self._cache_lock: + self._preparing.pop(key, None) + pending.error = error + pending.event.set() + raise + publication_committed = False + try: + with self._cache_lock: + old_misses = self._misses + try: + self._slots[key] = slot + self._misses = old_misses + 1 + self._preparing.pop(key, None) + pending.event.set() + publication_committed = True + except BaseException as error: + if self._slots.get(key) is slot: + self._slots.pop(key, None) + self._misses = old_misses + if self._preparing.get(key) is pending: + self._preparing.pop(key, None) + pending.error = error + pending.event.set() + raise + except BaseException as error: + if not publication_committed: + with self._cache_lock: + if self._slots.get(key) is slot: + self._slots.pop(key, None) + self._misses -= 1 + if self._preparing.get(key) is pending: + self._preparing.pop(key, None) + if pending.error is None: + pending.error = error + pending.event.set() + slot.lock.release() + raise + return slot, True + + def clear(self, *, synchronize: bool = True) -> None: + """Exclusively release slots after admitted host calls finish. + + The default waits for device completion and destroys the driver graph + handles eagerly. With ``synchronize=False``, every plan-held launch + argument, slot-owned norm buffer, and pooled default output is tied to + its plan's stream via ``record_stream`` before release, keeping tensor + storage allocator-safe; graph handles are only dropped (an executing + graph must never be destroyed underneath the device) and the caller + remains responsible for observing asynchronous completion. + """ + import torch + + with self._lifecycle: + while self._clearing: + self._lifecycle.wait() + try: + self._clearing = True + while self._active_computes: + self._lifecycle.wait() + if synchronize: + with torch.cuda.device(self.device_index): + torch.cuda.synchronize() + with self._cache_lock: + if synchronize: + for slot in self._slots.values(): + if slot.graph is not None: + slot.graph.destroy() + else: + for slot in self._slots.values(): + plan_stream = slot.plan.torch_stream + slot.plan.record_stream(plan_stream) + if slot.norm_plan is not None: + slot.norm_plan.record_stream(plan_stream) + _record_stream( + plan_stream, + slot.x_sq, + slot.c_sq, + *slot.default_outputs, + ) + self._slots.clear() + self._hits = 0 + self._misses = 0 + finally: + self._clearing = False + self._lifecycle.notify_all() + + def warmup(self, *args: Any, synchronize: bool = True, **kwargs: Any): + result = self.compute(*args, **kwargs) + if synchronize: + import torch + + with torch.cuda.device(self.device_index): + torch.cuda.synchronize() + return result + + +def init( + *, + device: Any = None, + arch: str | None = None, + timeout_ms: float | None = None, + max_cached_shapes: int | None = None, + compile: str = "lazy", +) -> FlashKMeansAssignRuntime: + """Create one reusable runtime; shape-specific plans remain lazy.""" + + return FlashKMeansAssignRuntime( + device=device, + arch=arch, + timeout_ms=timeout_ms, + max_cached_shapes=max_cached_shapes, + compile=compile, + ) + + +def _require_aux_tensor( + tensor: Any, + *, + name: str, + shape: tuple[int, ...], + dtype: Any, + device: Any, +) -> None: + if tuple(tensor.shape) != shape or tensor.dtype is not dtype: + raise ValueError(f"{name} must have dtype {dtype} and shape {shape}") + if tensor.device != device or not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous and on {device}") + + +def _record_stream(stream: Any, *tensors: Any) -> None: + seen: set[int] = set() + for tensor in tensors: + if tensor is None: + continue + identity = id(tensor) + if identity in seen: + continue + seen.add(identity) + record_stream = getattr(tensor, "record_stream", None) + if callable(record_stream): + record_stream(stream) + + +def _tensor_device_index(tensor: Any) -> int: + import torch + + index = tensor.device.index + return int(torch.cuda.current_device() if index is None else index) + + +def _prepare_inputs( + x: Any, + centroids: Any, + *, + out: Any | None, + x_sq: Any | None, + c_sq: Any | None, + device_index: int | None, + arch: str | None, + stream: Any, + arch_is_validated: bool = False, + timeout_ms: float | None = None, + defer_missing_norms: bool = False, + _context_is_active: bool = False, + _torch: Any = None, +) -> tuple[dict[str, Any], Any, str]: + if _torch is None: + import torch + else: + torch = _torch + + if not all(isinstance(item, torch.Tensor) and item.is_cuda for item in (x, centroids)): + raise TypeError("x and centroids must be CUDA torch.Tensor objects") + if x.dtype is not torch.bfloat16 or centroids.dtype is not torch.bfloat16: + raise TypeError("x and centroids dtype must be bfloat16") + if x.ndim != 3 or centroids.ndim != 3 or not x.is_contiguous() or not centroids.is_contiguous(): + raise ValueError("x and centroids must be contiguous [B, rows, D] tensors") + bsz, n_points, dim = map(int, x.shape) + c_bsz, n_clusters, c_dim = map(int, centroids.shape) + if (bsz, dim) != (c_bsz, c_dim) or x.device != centroids.device: + raise ValueError("x and centroids batch/feature dimensions and device must match") + input_device_index = x.device.index + if input_device_index is None: + input_device_index = torch.cuda.current_device() + input_device_index = int(input_device_index) + if device_index is not None and input_device_index != int(device_index): + raise ValueError( + f"Flash-KMeans runtime targets CUDA device {device_index}, got input device {input_device_index}" + ) + with _device_stream_context( + torch, + input_device_index, + stream, + context_is_active=_context_is_active, + ) as resolved_stream: + if arch_is_validated: + if arch is None: + raise RuntimeError("validated Flash-KMeans runtime must provide its active arch") + resolved_arch = str(arch) + else: + detected_arch = str(detect_gpu_arch()) + resolved_arch = detected_arch if arch is None else str(arch) + if resolved_arch != detected_arch: + raise ValueError( + "Flash-KMeans launch arch must match the active device: " + f"requested {resolved_arch}, detected {detected_arch}" + ) + caller_supplied_out = out is not None + if not caller_supplied_out: + out = torch.empty((bsz, n_points), dtype=torch.int32, device=x.device) + else: + _require_aux_tensor( + out, + name="out", + shape=(bsz, n_points), + dtype=torch.int32, + device=x.device, + ) + compute_x_sq = x_sq is None + compute_c_sq = c_sq is None + if not defer_missing_norms: + if compute_x_sq: + x_sq = torch.empty((bsz, n_points), dtype=torch.float32, device=x.device) + if compute_c_sq: + c_sq = torch.empty((bsz, n_clusters), dtype=torch.float32, device=x.device) + if x_sq is not None: + _require_aux_tensor( + x_sq, + name="x_sq", + shape=(bsz, n_points), + dtype=torch.float32, + device=x.device, + ) + if c_sq is not None: + _require_aux_tensor( + c_sq, + name="c_sq", + shape=(bsz, n_clusters), + dtype=torch.float32, + device=x.device, + ) + if (compute_x_sq or compute_c_sq) and not defer_missing_norms: + # This single custom launch converts BF16, squares, and reduces + # both row sets without PyTorch conversion/reduction temporaries. + from ._row_norm import launch_bf16_pair_row_norm + + _record_stream(resolved_stream, x, centroids, x_sq, c_sq, out) + launch_bf16_pair_row_norm( + x, + centroids, + x_sq, + c_sq, + compute_x=compute_x_sq, + compute_c=compute_c_sq, + stream=resolved_stream, + arch=resolved_arch, + timeout_ms=timeout_ms, + ) + norm_mode = ( + "fused_bf16_pair_row_norm" + if compute_x_sq and compute_c_sq + else "mixed_fused_bf16_pair_row_norm" + if compute_x_sq or compute_c_sq + else "explicit_precomputed" + ) + inputs = { + "B": bsz, + "N": n_points, + "D": dim, + "K": n_clusters, + "dtype": "bfloat16", + "x": x, + "centroids": centroids, + "x_sq": x_sq, + "c_sq": c_sq, + "out": out, + "_norm_mode": norm_mode, + } + return inputs, resolved_stream, resolved_arch + + +def prepare_flash_kmeans_assign( + x: Any, + centroids: Any, + *, + out: Any | None = None, + x_sq: Any | None = None, + c_sq: Any | None = None, + arch: str | None = None, + stream: Any = None, + timeout_ms: float | None = None, +) -> PreparedFlashKMeansAssign: + """Validate and prepare an allocation-free direct launch for exact tensors.""" + + timeout_ms = _validate_timeout(timeout_ms) + + inputs, resolved_stream, resolved_arch = _prepare_inputs( + x, + centroids, + out=out, + x_sq=x_sq, + c_sq=c_sq, + device_index=None, + arch=arch, + stream=stream, + timeout_ms=timeout_ms, + ) + import torch + + input_device_index = inputs["x"].device.index + device_index = int(torch.cuda.current_device() if input_device_index is None else input_device_index) + with torch.cuda.device(device_index), torch.cuda.stream(resolved_stream): + launch_plan = _dispatcher.prepare_launch_plan( + inputs, + arch=resolved_arch, + stream=resolved_stream, + timeout_ms=timeout_ms, + ) + return PreparedFlashKMeansAssign( + inputs=inputs, + launch_plan=launch_plan, + stream=resolved_stream, + timeout_ms=timeout_ms, + ) + + +def flash_kmeans_assign_prepared( + prepared: PreparedFlashKMeansAssign, + *, + stream: Any = None, + timeout_ms: float | None = None, + return_info: bool = False, +): + """Submit a prepared plan without route selection, allocation, or packing.""" + + if not isinstance(prepared, PreparedFlashKMeansAssign): + raise TypeError("prepared must be returned by prepare_flash_kmeans_assign") + effective_timeout_ms = prepared.timeout_ms if timeout_ms is None else _validate_timeout(timeout_ms) + resolved_stream = prepared.stream if stream is None else stream + _record_stream( + resolved_stream, + *(prepared.inputs[name] for name in ("x", "centroids", "x_sq", "c_sq", "out")), + ) + result = _dispatcher.launch_prepared( + prepared.launch_plan, + stream=resolved_stream, + timeout_ms=effective_timeout_ms, + public_inputs_already_recorded=True, + ) + produced = result.get("cluster_ids", prepared.out) if isinstance(result, dict) else prepared.out + if produced is not prepared.out: + prepared.out.copy_(produced) + if not return_info: + return prepared.out + info = { + "semantic_entrypoint": SEMANTIC_ENTRYPOINT, + "selected_route": prepared.selected_route, + "launch_entrypoint": prepared.launch_plan.launch_entrypoint, + "exact_launch_plan": True, + "prepared_launch_count": prepared.launch_plan.launch_count, + "arch": prepared.launch_plan.arch, + "device_index": prepared.launch_plan.device_index, + "stream_handle": prepared.launch_plan.stream_handle, + } + return prepared.out, info + + +_DEFAULT_RUNTIME_LOCK = RLock() +_DEFAULT_RUNTIMES: dict[int, FlashKMeansAssignRuntime] = {} + + +def _default_runtime(device_index: int) -> FlashKMeansAssignRuntime: + """Return the process-wide per-device runtime backing ``flash_kmeans_assign``.""" + + runtime = _DEFAULT_RUNTIMES.get(device_index) + if runtime is None: + with _DEFAULT_RUNTIME_LOCK: + runtime = _DEFAULT_RUNTIMES.get(device_index) + if runtime is None: + runtime = FlashKMeansAssignRuntime(device=device_index) + _DEFAULT_RUNTIMES[device_index] = runtime + return runtime + + +def flash_kmeans_assign( + x: Any, + centroids: Any, + *, + out: Any | None = None, + x_sq: Any | None = None, + c_sq: Any | None = None, + arch: str | None = None, + stream: Any = None, + timeout_ms: float | None = None, + return_info: bool = False, +): + """Assign points to centroids through the per-device signature plan cache. + + The first call for a signature runs the frozen guard cascade once to + construct its launch plan and CUDA graph; subsequent calls are + pointer-overwrite graph replays on the plan's stream. A caller-provided + ``out=`` is written through; default outputs come from the signature's + plan-owned pool, so a default-output result must be consumed (or cloned) + before two further default-output calls of the same signature overwrite + its storage. + """ + import torch + + if not isinstance(x, torch.Tensor) or not getattr(x, "is_cuda", False): + raise TypeError("x must be a CUDA torch.Tensor") + runtime = _default_runtime(_tensor_device_index(x)) + if arch is not None and str(arch) != runtime.arch: + raise ValueError( + f"Flash-KMeans launch arch must match the active device: requested {arch}, detected {runtime.arch}" + ) + return runtime.compute( + x, + centroids, + out=out, + x_sq=x_sq, + c_sq=c_sq, + stream=stream, + timeout_ms=timeout_ms, + return_info=return_info, + ) diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/kernels.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/kernels.py new file mode 100644 index 00000000..32b33d0f --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/kernels.py @@ -0,0 +1,9150 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass +from importlib import resources +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from ._runtime import CUDAKernel + + +@dataclass(frozen=True) +class KernelSpec: + name: str + symbol: str + source: str + threads: int + shared_mem_bytes: int + cluster_dims: tuple[int, int, int] + launch_mode: str + parameters: tuple[dict[str, str], ...] + specializations: dict[str, int | str] + compile_options: tuple[str, ...] + + @staticmethod + def from_manifest(entry: dict[str, Any]) -> "KernelSpec": + return KernelSpec( + name=entry["name"], + symbol=entry["symbol"], + source=entry["source"], + threads=int(entry["threads"]), + shared_mem_bytes=int(entry["shared_mem_bytes"]), + cluster_dims=tuple(int(v) for v in entry["cluster_dims"]), + launch_mode=entry["launch_mode"], + parameters=tuple(entry["parameters"]), + specializations=dict(entry.get("specializations", {})), + compile_options=tuple(str(option) for option in entry.get("compile_options", ())), + ) + + +class ExportedKernel: + def __init__(self, spec: KernelSpec): + self.spec = spec + self._compiled: dict[tuple[str, tuple[str, ...], int, int], "CUDAKernel"] = {} + self._arg_types = tuple(parameter["ctype"] for parameter in spec.parameters) + self._default_block = (spec.threads, 1, 1) + self._default_shared_mem = spec.shared_mem_bytes + self._source_cache: str | None = None + self._source_digest: str | None = None + + @property + def parameters(self) -> tuple[dict[str, str], ...]: + return self.spec.parameters + + @property + def arg_types(self) -> tuple[str, ...]: + return self._arg_types + + def source_text(self) -> str: + if self._source_cache is None: + from ._runtime import record_source_read + + package = __package__ or __name__.rpartition(".")[0] + self._source_cache = resources.files(package).joinpath(self.spec.source).read_text(encoding="utf-8") + record_source_read() + return self._source_cache + + def compile(self, *, arch: str | None = None, options: list[str] | None = None) -> "CUDAKernel": + effective_options = tuple(dict.fromkeys((*self.spec.compile_options, *(options or ())))) + from ._runtime import ( + compilation_cache_generation, + current_cuda_device_index, + load_cached_kernel, + resolve_gpu_arch, + ) + + resolved_arch = resolve_gpu_arch(arch) + device_index = current_cuda_device_index() + generation = compilation_cache_generation() + key = (resolved_arch, effective_options, device_index, generation) + kernel = self._compiled.get(key) + if kernel is None or kernel.closed: + source = self.source_text() + if self._source_digest is None: + import hashlib + + self._source_digest = hashlib.sha256(source.encode("utf-8")).hexdigest() + kernel = load_cached_kernel( + source, + source_digest=self._source_digest, + func_name=self.spec.symbol, + arch=resolved_arch, + device_index=device_index, + name=f"{self.spec.name}.cu", + options=effective_options, + ) + self._compiled[key] = kernel + return kernel + + def launch( + self, + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, + ) -> None: + from ._runtime import resolve_launch_defaults + + arch, stream, timeout_ms = resolve_launch_defaults( + arch=arch, + stream=stream, + timeout_ms=timeout_ms, + ) + prepared = self.prepare_launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + arch=arch, + options=options, + ) + prepared.launch(timeout_ms=timeout_ms) + + def prepare_launch( + self, + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + arch: str | None = None, + options: list[str] | None = None, + ): + # Compile and fully marshal one launch without submitting GPU work. + + if len(args) != len(self.spec.parameters): + expected = ", ".join(p["name"] for p in self.spec.parameters) + raise TypeError(f"{self.spec.name} expects {len(self.spec.parameters)} args ({expected}), got {len(args)}") + from ._runtime import launch_stream_context, resolve_launch_defaults + + arch, stream, _ = resolve_launch_defaults( + arch=arch, + stream=stream, + timeout_ms=None, + ) + if block is None: + block = self._default_block + if shared_mem is None: + shared_mem = self._default_shared_mem + with launch_stream_context(stream): + kernel = self.compile(arch=arch, options=options) + if self.spec.launch_mode == "cluster": + return kernel.prepare_launch_cluster( + grid=grid, + block=block, + args=args, + arg_types=self._arg_types, + cluster_dims=self.spec.cluster_dims, + shared_mem=shared_mem, + stream=stream, + ) + if self.spec.launch_mode == "cooperative": + return kernel.prepare_launch_cooperative( + grid=grid, + block=block, + args=args, + arg_types=self._arg_types, + shared_mem=shared_mem, + stream=stream, + ) + return kernel.prepare_launch( + grid=grid, + block=block, + args=args, + arg_types=self._arg_types, + shared_mem=shared_mem, + stream=stream, + ) + + +def _load_manifest() -> dict[str, Any]: + package = __package__ or __name__.rpartition(".")[0] + return json.loads(resources.files(package).joinpath("manifest.json").read_text(encoding="utf-8")) + + +_MANIFEST = _load_manifest() +KERNELS = {entry["name"]: ExportedKernel(KernelSpec.from_manifest(entry)) for entry in _MANIFEST["kernels"]} + + +def get_kernel(name: str) -> ExportedKernel: + try: + return KERNELS[name] + except KeyError as exc: + available = ", ".join(sorted(KERNELS)) + raise KeyError(f"Unknown exported kernel {name!r}. Available: {available}") from exc + + +dispatch_kernel_0000 = get_kernel('dispatch_kernel_0000') + + +def launch_dispatch_kernel_0000( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0000.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0001 = get_kernel('dispatch_kernel_0001') + + +def launch_dispatch_kernel_0001( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0001.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0002 = get_kernel('dispatch_kernel_0002') + + +def launch_dispatch_kernel_0002( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0002.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0003 = get_kernel('dispatch_kernel_0003') + + +def launch_dispatch_kernel_0003( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0003.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0004 = get_kernel('dispatch_kernel_0004') + + +def launch_dispatch_kernel_0004( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0004.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0005 = get_kernel('dispatch_kernel_0005') + + +def launch_dispatch_kernel_0005( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0005.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0006 = get_kernel('dispatch_kernel_0006') + + +def launch_dispatch_kernel_0006( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0006.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0007 = get_kernel('dispatch_kernel_0007') + + +def launch_dispatch_kernel_0007( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0007.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0008 = get_kernel('dispatch_kernel_0008') + + +def launch_dispatch_kernel_0008( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0008.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0009 = get_kernel('dispatch_kernel_0009') + + +def launch_dispatch_kernel_0009( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0009.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0010 = get_kernel('dispatch_kernel_0010') + + +def launch_dispatch_kernel_0010( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0010.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0011 = get_kernel('dispatch_kernel_0011') + + +def launch_dispatch_kernel_0011( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0011.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0012 = get_kernel('dispatch_kernel_0012') + + +def launch_dispatch_kernel_0012( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0012.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0013 = get_kernel('dispatch_kernel_0013') + + +def launch_dispatch_kernel_0013( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0013.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0014 = get_kernel('dispatch_kernel_0014') + + +def launch_dispatch_kernel_0014( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0014.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0015 = get_kernel('dispatch_kernel_0015') + + +def launch_dispatch_kernel_0015( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0015.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0016 = get_kernel('dispatch_kernel_0016') + + +def launch_dispatch_kernel_0016( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0016.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0017 = get_kernel('dispatch_kernel_0017') + + +def launch_dispatch_kernel_0017( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0017.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0018 = get_kernel('dispatch_kernel_0018') + + +def launch_dispatch_kernel_0018( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0018.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0019 = get_kernel('dispatch_kernel_0019') + + +def launch_dispatch_kernel_0019( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0019.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0020 = get_kernel('dispatch_kernel_0020') + + +def launch_dispatch_kernel_0020( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0020.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0021 = get_kernel('dispatch_kernel_0021') + + +def launch_dispatch_kernel_0021( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0021.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0022 = get_kernel('dispatch_kernel_0022') + + +def launch_dispatch_kernel_0022( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0022.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0023 = get_kernel('dispatch_kernel_0023') + + +def launch_dispatch_kernel_0023( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0023.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0024 = get_kernel('dispatch_kernel_0024') + + +def launch_dispatch_kernel_0024( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0024.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0025 = get_kernel('dispatch_kernel_0025') + + +def launch_dispatch_kernel_0025( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0025.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0026 = get_kernel('dispatch_kernel_0026') + + +def launch_dispatch_kernel_0026( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0026.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0027 = get_kernel('dispatch_kernel_0027') + + +def launch_dispatch_kernel_0027( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0027.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0028 = get_kernel('dispatch_kernel_0028') + + +def launch_dispatch_kernel_0028( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0028.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0029 = get_kernel('dispatch_kernel_0029') + + +def launch_dispatch_kernel_0029( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0029.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0030 = get_kernel('dispatch_kernel_0030') + + +def launch_dispatch_kernel_0030( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0030.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0031 = get_kernel('dispatch_kernel_0031') + + +def launch_dispatch_kernel_0031( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0031.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0032 = get_kernel('dispatch_kernel_0032') + + +def launch_dispatch_kernel_0032( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0032.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0033 = get_kernel('dispatch_kernel_0033') + + +def launch_dispatch_kernel_0033( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0033.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0034 = get_kernel('dispatch_kernel_0034') + + +def launch_dispatch_kernel_0034( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0034.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0035 = get_kernel('dispatch_kernel_0035') + + +def launch_dispatch_kernel_0035( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0035.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0036 = get_kernel('dispatch_kernel_0036') + + +def launch_dispatch_kernel_0036( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0036.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0037 = get_kernel('dispatch_kernel_0037') + + +def launch_dispatch_kernel_0037( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0037.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0038 = get_kernel('dispatch_kernel_0038') + + +def launch_dispatch_kernel_0038( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0038.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0039 = get_kernel('dispatch_kernel_0039') + + +def launch_dispatch_kernel_0039( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0039.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0040 = get_kernel('dispatch_kernel_0040') + + +def launch_dispatch_kernel_0040( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0040.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0041 = get_kernel('dispatch_kernel_0041') + + +def launch_dispatch_kernel_0041( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0041.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0042 = get_kernel('dispatch_kernel_0042') + + +def launch_dispatch_kernel_0042( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0042.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0043 = get_kernel('dispatch_kernel_0043') + + +def launch_dispatch_kernel_0043( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0043.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0044 = get_kernel('dispatch_kernel_0044') + + +def launch_dispatch_kernel_0044( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0044.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0045 = get_kernel('dispatch_kernel_0045') + + +def launch_dispatch_kernel_0045( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0045.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0046 = get_kernel('dispatch_kernel_0046') + + +def launch_dispatch_kernel_0046( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0046.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0047 = get_kernel('dispatch_kernel_0047') + + +def launch_dispatch_kernel_0047( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0047.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0048 = get_kernel('dispatch_kernel_0048') + + +def launch_dispatch_kernel_0048( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0048.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0049 = get_kernel('dispatch_kernel_0049') + + +def launch_dispatch_kernel_0049( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0049.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0050 = get_kernel('dispatch_kernel_0050') + + +def launch_dispatch_kernel_0050( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0050.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0051 = get_kernel('dispatch_kernel_0051') + + +def launch_dispatch_kernel_0051( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0051.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0052 = get_kernel('dispatch_kernel_0052') + + +def launch_dispatch_kernel_0052( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0052.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0053 = get_kernel('dispatch_kernel_0053') + + +def launch_dispatch_kernel_0053( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0053.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0054 = get_kernel('dispatch_kernel_0054') + + +def launch_dispatch_kernel_0054( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0054.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0055 = get_kernel('dispatch_kernel_0055') + + +def launch_dispatch_kernel_0055( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0055.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0056 = get_kernel('dispatch_kernel_0056') + + +def launch_dispatch_kernel_0056( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0056.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0057 = get_kernel('dispatch_kernel_0057') + + +def launch_dispatch_kernel_0057( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0057.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0058 = get_kernel('dispatch_kernel_0058') + + +def launch_dispatch_kernel_0058( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0058.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0059 = get_kernel('dispatch_kernel_0059') + + +def launch_dispatch_kernel_0059( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0059.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0060 = get_kernel('dispatch_kernel_0060') + + +def launch_dispatch_kernel_0060( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0060.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0061 = get_kernel('dispatch_kernel_0061') + + +def launch_dispatch_kernel_0061( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0061.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0062 = get_kernel('dispatch_kernel_0062') + + +def launch_dispatch_kernel_0062( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0062.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0063 = get_kernel('dispatch_kernel_0063') + + +def launch_dispatch_kernel_0063( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0063.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0064 = get_kernel('dispatch_kernel_0064') + + +def launch_dispatch_kernel_0064( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0064.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0065 = get_kernel('dispatch_kernel_0065') + + +def launch_dispatch_kernel_0065( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0065.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0066 = get_kernel('dispatch_kernel_0066') + + +def launch_dispatch_kernel_0066( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0066.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0067 = get_kernel('dispatch_kernel_0067') + + +def launch_dispatch_kernel_0067( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0067.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0068 = get_kernel('dispatch_kernel_0068') + + +def launch_dispatch_kernel_0068( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0068.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0069 = get_kernel('dispatch_kernel_0069') + + +def launch_dispatch_kernel_0069( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0069.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0070 = get_kernel('dispatch_kernel_0070') + + +def launch_dispatch_kernel_0070( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0070.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0071 = get_kernel('dispatch_kernel_0071') + + +def launch_dispatch_kernel_0071( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0071.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0072 = get_kernel('dispatch_kernel_0072') + + +def launch_dispatch_kernel_0072( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0072.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0073 = get_kernel('dispatch_kernel_0073') + + +def launch_dispatch_kernel_0073( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0073.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0074 = get_kernel('dispatch_kernel_0074') + + +def launch_dispatch_kernel_0074( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0074.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0075 = get_kernel('dispatch_kernel_0075') + + +def launch_dispatch_kernel_0075( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0075.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0076 = get_kernel('dispatch_kernel_0076') + + +def launch_dispatch_kernel_0076( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0076.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0077 = get_kernel('dispatch_kernel_0077') + + +def launch_dispatch_kernel_0077( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0077.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0078 = get_kernel('dispatch_kernel_0078') + + +def launch_dispatch_kernel_0078( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0078.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0079 = get_kernel('dispatch_kernel_0079') + + +def launch_dispatch_kernel_0079( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0079.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0080 = get_kernel('dispatch_kernel_0080') + + +def launch_dispatch_kernel_0080( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0080.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0081 = get_kernel('dispatch_kernel_0081') + + +def launch_dispatch_kernel_0081( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0081.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0082 = get_kernel('dispatch_kernel_0082') + + +def launch_dispatch_kernel_0082( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0082.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0083 = get_kernel('dispatch_kernel_0083') + + +def launch_dispatch_kernel_0083( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0083.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0084 = get_kernel('dispatch_kernel_0084') + + +def launch_dispatch_kernel_0084( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0084.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0085 = get_kernel('dispatch_kernel_0085') + + +def launch_dispatch_kernel_0085( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0085.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0086 = get_kernel('dispatch_kernel_0086') + + +def launch_dispatch_kernel_0086( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0086.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0087 = get_kernel('dispatch_kernel_0087') + + +def launch_dispatch_kernel_0087( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0087.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0088 = get_kernel('dispatch_kernel_0088') + + +def launch_dispatch_kernel_0088( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0088.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0089 = get_kernel('dispatch_kernel_0089') + + +def launch_dispatch_kernel_0089( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0089.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0090 = get_kernel('dispatch_kernel_0090') + + +def launch_dispatch_kernel_0090( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0090.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0091 = get_kernel('dispatch_kernel_0091') + + +def launch_dispatch_kernel_0091( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0091.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0092 = get_kernel('dispatch_kernel_0092') + + +def launch_dispatch_kernel_0092( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0092.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0093 = get_kernel('dispatch_kernel_0093') + + +def launch_dispatch_kernel_0093( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0093.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0094 = get_kernel('dispatch_kernel_0094') + + +def launch_dispatch_kernel_0094( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0094.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0095 = get_kernel('dispatch_kernel_0095') + + +def launch_dispatch_kernel_0095( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0095.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0096 = get_kernel('dispatch_kernel_0096') + + +def launch_dispatch_kernel_0096( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0096.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0097 = get_kernel('dispatch_kernel_0097') + + +def launch_dispatch_kernel_0097( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0097.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0098 = get_kernel('dispatch_kernel_0098') + + +def launch_dispatch_kernel_0098( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0098.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0099 = get_kernel('dispatch_kernel_0099') + + +def launch_dispatch_kernel_0099( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0099.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0100 = get_kernel('dispatch_kernel_0100') + + +def launch_dispatch_kernel_0100( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0100.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0101 = get_kernel('dispatch_kernel_0101') + + +def launch_dispatch_kernel_0101( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0101.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0102 = get_kernel('dispatch_kernel_0102') + + +def launch_dispatch_kernel_0102( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0102.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0103 = get_kernel('dispatch_kernel_0103') + + +def launch_dispatch_kernel_0103( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0103.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0104 = get_kernel('dispatch_kernel_0104') + + +def launch_dispatch_kernel_0104( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0104.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0105 = get_kernel('dispatch_kernel_0105') + + +def launch_dispatch_kernel_0105( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0105.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0106 = get_kernel('dispatch_kernel_0106') + + +def launch_dispatch_kernel_0106( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0106.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0107 = get_kernel('dispatch_kernel_0107') + + +def launch_dispatch_kernel_0107( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0107.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0108 = get_kernel('dispatch_kernel_0108') + + +def launch_dispatch_kernel_0108( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0108.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0109 = get_kernel('dispatch_kernel_0109') + + +def launch_dispatch_kernel_0109( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0109.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0110 = get_kernel('dispatch_kernel_0110') + + +def launch_dispatch_kernel_0110( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0110.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0111 = get_kernel('dispatch_kernel_0111') + + +def launch_dispatch_kernel_0111( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0111.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0112 = get_kernel('dispatch_kernel_0112') + + +def launch_dispatch_kernel_0112( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0112.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0113 = get_kernel('dispatch_kernel_0113') + + +def launch_dispatch_kernel_0113( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0113.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0114 = get_kernel('dispatch_kernel_0114') + + +def launch_dispatch_kernel_0114( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0114.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0115 = get_kernel('dispatch_kernel_0115') + + +def launch_dispatch_kernel_0115( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0115.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0116 = get_kernel('dispatch_kernel_0116') + + +def launch_dispatch_kernel_0116( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0116.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0117 = get_kernel('dispatch_kernel_0117') + + +def launch_dispatch_kernel_0117( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0117.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0118 = get_kernel('dispatch_kernel_0118') + + +def launch_dispatch_kernel_0118( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0118.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0119 = get_kernel('dispatch_kernel_0119') + + +def launch_dispatch_kernel_0119( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0119.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0120 = get_kernel('dispatch_kernel_0120') + + +def launch_dispatch_kernel_0120( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0120.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0121 = get_kernel('dispatch_kernel_0121') + + +def launch_dispatch_kernel_0121( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0121.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0122 = get_kernel('dispatch_kernel_0122') + + +def launch_dispatch_kernel_0122( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0122.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0123 = get_kernel('dispatch_kernel_0123') + + +def launch_dispatch_kernel_0123( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0123.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0124 = get_kernel('dispatch_kernel_0124') + + +def launch_dispatch_kernel_0124( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0124.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0125 = get_kernel('dispatch_kernel_0125') + + +def launch_dispatch_kernel_0125( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0125.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0126 = get_kernel('dispatch_kernel_0126') + + +def launch_dispatch_kernel_0126( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0126.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0127 = get_kernel('dispatch_kernel_0127') + + +def launch_dispatch_kernel_0127( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0127.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0128 = get_kernel('dispatch_kernel_0128') + + +def launch_dispatch_kernel_0128( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0128.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0129 = get_kernel('dispatch_kernel_0129') + + +def launch_dispatch_kernel_0129( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0129.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0130 = get_kernel('dispatch_kernel_0130') + + +def launch_dispatch_kernel_0130( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0130.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0131 = get_kernel('dispatch_kernel_0131') + + +def launch_dispatch_kernel_0131( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0131.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0132 = get_kernel('dispatch_kernel_0132') + + +def launch_dispatch_kernel_0132( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0132.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0133 = get_kernel('dispatch_kernel_0133') + + +def launch_dispatch_kernel_0133( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0133.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0134 = get_kernel('dispatch_kernel_0134') + + +def launch_dispatch_kernel_0134( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0134.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0135 = get_kernel('dispatch_kernel_0135') + + +def launch_dispatch_kernel_0135( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0135.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0136 = get_kernel('dispatch_kernel_0136') + + +def launch_dispatch_kernel_0136( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0136.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0137 = get_kernel('dispatch_kernel_0137') + + +def launch_dispatch_kernel_0137( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0137.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0138 = get_kernel('dispatch_kernel_0138') + + +def launch_dispatch_kernel_0138( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0138.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0139 = get_kernel('dispatch_kernel_0139') + + +def launch_dispatch_kernel_0139( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0139.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0140 = get_kernel('dispatch_kernel_0140') + + +def launch_dispatch_kernel_0140( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0140.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0141 = get_kernel('dispatch_kernel_0141') + + +def launch_dispatch_kernel_0141( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0141.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0142 = get_kernel('dispatch_kernel_0142') + + +def launch_dispatch_kernel_0142( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0142.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0143 = get_kernel('dispatch_kernel_0143') + + +def launch_dispatch_kernel_0143( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0143.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0144 = get_kernel('dispatch_kernel_0144') + + +def launch_dispatch_kernel_0144( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0144.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0145 = get_kernel('dispatch_kernel_0145') + + +def launch_dispatch_kernel_0145( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0145.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0146 = get_kernel('dispatch_kernel_0146') + + +def launch_dispatch_kernel_0146( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0146.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0147 = get_kernel('dispatch_kernel_0147') + + +def launch_dispatch_kernel_0147( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0147.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0148 = get_kernel('dispatch_kernel_0148') + + +def launch_dispatch_kernel_0148( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0148.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0149 = get_kernel('dispatch_kernel_0149') + + +def launch_dispatch_kernel_0149( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0149.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0150 = get_kernel('dispatch_kernel_0150') + + +def launch_dispatch_kernel_0150( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0150.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0151 = get_kernel('dispatch_kernel_0151') + + +def launch_dispatch_kernel_0151( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0151.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0152 = get_kernel('dispatch_kernel_0152') + + +def launch_dispatch_kernel_0152( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0152.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0153 = get_kernel('dispatch_kernel_0153') + + +def launch_dispatch_kernel_0153( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0153.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0154 = get_kernel('dispatch_kernel_0154') + + +def launch_dispatch_kernel_0154( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0154.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0155 = get_kernel('dispatch_kernel_0155') + + +def launch_dispatch_kernel_0155( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0155.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0156 = get_kernel('dispatch_kernel_0156') + + +def launch_dispatch_kernel_0156( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0156.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0157 = get_kernel('dispatch_kernel_0157') + + +def launch_dispatch_kernel_0157( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0157.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0158 = get_kernel('dispatch_kernel_0158') + + +def launch_dispatch_kernel_0158( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0158.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0159 = get_kernel('dispatch_kernel_0159') + + +def launch_dispatch_kernel_0159( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0159.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0160 = get_kernel('dispatch_kernel_0160') + + +def launch_dispatch_kernel_0160( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0160.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0161 = get_kernel('dispatch_kernel_0161') + + +def launch_dispatch_kernel_0161( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0161.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0162 = get_kernel('dispatch_kernel_0162') + + +def launch_dispatch_kernel_0162( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0162.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0163 = get_kernel('dispatch_kernel_0163') + + +def launch_dispatch_kernel_0163( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0163.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0164 = get_kernel('dispatch_kernel_0164') + + +def launch_dispatch_kernel_0164( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0164.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0165 = get_kernel('dispatch_kernel_0165') + + +def launch_dispatch_kernel_0165( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0165.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0166 = get_kernel('dispatch_kernel_0166') + + +def launch_dispatch_kernel_0166( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0166.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0167 = get_kernel('dispatch_kernel_0167') + + +def launch_dispatch_kernel_0167( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0167.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0168 = get_kernel('dispatch_kernel_0168') + + +def launch_dispatch_kernel_0168( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0168.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0169 = get_kernel('dispatch_kernel_0169') + + +def launch_dispatch_kernel_0169( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0169.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0170 = get_kernel('dispatch_kernel_0170') + + +def launch_dispatch_kernel_0170( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0170.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0171 = get_kernel('dispatch_kernel_0171') + + +def launch_dispatch_kernel_0171( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0171.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0172 = get_kernel('dispatch_kernel_0172') + + +def launch_dispatch_kernel_0172( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0172.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0173 = get_kernel('dispatch_kernel_0173') + + +def launch_dispatch_kernel_0173( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0173.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0174 = get_kernel('dispatch_kernel_0174') + + +def launch_dispatch_kernel_0174( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0174.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0175 = get_kernel('dispatch_kernel_0175') + + +def launch_dispatch_kernel_0175( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0175.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0176 = get_kernel('dispatch_kernel_0176') + + +def launch_dispatch_kernel_0176( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0176.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0177 = get_kernel('dispatch_kernel_0177') + + +def launch_dispatch_kernel_0177( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0177.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0178 = get_kernel('dispatch_kernel_0178') + + +def launch_dispatch_kernel_0178( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0178.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0179 = get_kernel('dispatch_kernel_0179') + + +def launch_dispatch_kernel_0179( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0179.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0180 = get_kernel('dispatch_kernel_0180') + + +def launch_dispatch_kernel_0180( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0180.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0181 = get_kernel('dispatch_kernel_0181') + + +def launch_dispatch_kernel_0181( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0181.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0182 = get_kernel('dispatch_kernel_0182') + + +def launch_dispatch_kernel_0182( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0182.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0183 = get_kernel('dispatch_kernel_0183') + + +def launch_dispatch_kernel_0183( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0183.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0184 = get_kernel('dispatch_kernel_0184') + + +def launch_dispatch_kernel_0184( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0184.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0185 = get_kernel('dispatch_kernel_0185') + + +def launch_dispatch_kernel_0185( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0185.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0186 = get_kernel('dispatch_kernel_0186') + + +def launch_dispatch_kernel_0186( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0186.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0187 = get_kernel('dispatch_kernel_0187') + + +def launch_dispatch_kernel_0187( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0187.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0188 = get_kernel('dispatch_kernel_0188') + + +def launch_dispatch_kernel_0188( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0188.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0189 = get_kernel('dispatch_kernel_0189') + + +def launch_dispatch_kernel_0189( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0189.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0190 = get_kernel('dispatch_kernel_0190') + + +def launch_dispatch_kernel_0190( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0190.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0191 = get_kernel('dispatch_kernel_0191') + + +def launch_dispatch_kernel_0191( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0191.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0192 = get_kernel('dispatch_kernel_0192') + + +def launch_dispatch_kernel_0192( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0192.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0193 = get_kernel('dispatch_kernel_0193') + + +def launch_dispatch_kernel_0193( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0193.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0194 = get_kernel('dispatch_kernel_0194') + + +def launch_dispatch_kernel_0194( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0194.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0195 = get_kernel('dispatch_kernel_0195') + + +def launch_dispatch_kernel_0195( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0195.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0196 = get_kernel('dispatch_kernel_0196') + + +def launch_dispatch_kernel_0196( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0196.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0197 = get_kernel('dispatch_kernel_0197') + + +def launch_dispatch_kernel_0197( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0197.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0198 = get_kernel('dispatch_kernel_0198') + + +def launch_dispatch_kernel_0198( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0198.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0199 = get_kernel('dispatch_kernel_0199') + + +def launch_dispatch_kernel_0199( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0199.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0200 = get_kernel('dispatch_kernel_0200') + + +def launch_dispatch_kernel_0200( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0200.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0201 = get_kernel('dispatch_kernel_0201') + + +def launch_dispatch_kernel_0201( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0201.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0202 = get_kernel('dispatch_kernel_0202') + + +def launch_dispatch_kernel_0202( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0202.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0203 = get_kernel('dispatch_kernel_0203') + + +def launch_dispatch_kernel_0203( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0203.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0204 = get_kernel('dispatch_kernel_0204') + + +def launch_dispatch_kernel_0204( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0204.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0205 = get_kernel('dispatch_kernel_0205') + + +def launch_dispatch_kernel_0205( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0205.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0206 = get_kernel('dispatch_kernel_0206') + + +def launch_dispatch_kernel_0206( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0206.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0207 = get_kernel('dispatch_kernel_0207') + + +def launch_dispatch_kernel_0207( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0207.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0208 = get_kernel('dispatch_kernel_0208') + + +def launch_dispatch_kernel_0208( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0208.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0209 = get_kernel('dispatch_kernel_0209') + + +def launch_dispatch_kernel_0209( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0209.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0210 = get_kernel('dispatch_kernel_0210') + + +def launch_dispatch_kernel_0210( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0210.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0211 = get_kernel('dispatch_kernel_0211') + + +def launch_dispatch_kernel_0211( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0211.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0212 = get_kernel('dispatch_kernel_0212') + + +def launch_dispatch_kernel_0212( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0212.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0213 = get_kernel('dispatch_kernel_0213') + + +def launch_dispatch_kernel_0213( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0213.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0214 = get_kernel('dispatch_kernel_0214') + + +def launch_dispatch_kernel_0214( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0214.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0215 = get_kernel('dispatch_kernel_0215') + + +def launch_dispatch_kernel_0215( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0215.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0216 = get_kernel('dispatch_kernel_0216') + + +def launch_dispatch_kernel_0216( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0216.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0217 = get_kernel('dispatch_kernel_0217') + + +def launch_dispatch_kernel_0217( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0217.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0218 = get_kernel('dispatch_kernel_0218') + + +def launch_dispatch_kernel_0218( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0218.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0219 = get_kernel('dispatch_kernel_0219') + + +def launch_dispatch_kernel_0219( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0219.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0220 = get_kernel('dispatch_kernel_0220') + + +def launch_dispatch_kernel_0220( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0220.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0221 = get_kernel('dispatch_kernel_0221') + + +def launch_dispatch_kernel_0221( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0221.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0222 = get_kernel('dispatch_kernel_0222') + + +def launch_dispatch_kernel_0222( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0222.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0223 = get_kernel('dispatch_kernel_0223') + + +def launch_dispatch_kernel_0223( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0223.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0224 = get_kernel('dispatch_kernel_0224') + + +def launch_dispatch_kernel_0224( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0224.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0225 = get_kernel('dispatch_kernel_0225') + + +def launch_dispatch_kernel_0225( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0225.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0226 = get_kernel('dispatch_kernel_0226') + + +def launch_dispatch_kernel_0226( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0226.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0227 = get_kernel('dispatch_kernel_0227') + + +def launch_dispatch_kernel_0227( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0227.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0228 = get_kernel('dispatch_kernel_0228') + + +def launch_dispatch_kernel_0228( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0228.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0229 = get_kernel('dispatch_kernel_0229') + + +def launch_dispatch_kernel_0229( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0229.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0230 = get_kernel('dispatch_kernel_0230') + + +def launch_dispatch_kernel_0230( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0230.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0231 = get_kernel('dispatch_kernel_0231') + + +def launch_dispatch_kernel_0231( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0231.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0232 = get_kernel('dispatch_kernel_0232') + + +def launch_dispatch_kernel_0232( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0232.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0233 = get_kernel('dispatch_kernel_0233') + + +def launch_dispatch_kernel_0233( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0233.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0234 = get_kernel('dispatch_kernel_0234') + + +def launch_dispatch_kernel_0234( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0234.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0235 = get_kernel('dispatch_kernel_0235') + + +def launch_dispatch_kernel_0235( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0235.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0236 = get_kernel('dispatch_kernel_0236') + + +def launch_dispatch_kernel_0236( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0236.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0237 = get_kernel('dispatch_kernel_0237') + + +def launch_dispatch_kernel_0237( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0237.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0238 = get_kernel('dispatch_kernel_0238') + + +def launch_dispatch_kernel_0238( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0238.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0239 = get_kernel('dispatch_kernel_0239') + + +def launch_dispatch_kernel_0239( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0239.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0240 = get_kernel('dispatch_kernel_0240') + + +def launch_dispatch_kernel_0240( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0240.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0241 = get_kernel('dispatch_kernel_0241') + + +def launch_dispatch_kernel_0241( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0241.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0242 = get_kernel('dispatch_kernel_0242') + + +def launch_dispatch_kernel_0242( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0242.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0243 = get_kernel('dispatch_kernel_0243') + + +def launch_dispatch_kernel_0243( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0243.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0244 = get_kernel('dispatch_kernel_0244') + + +def launch_dispatch_kernel_0244( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0244.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0245 = get_kernel('dispatch_kernel_0245') + + +def launch_dispatch_kernel_0245( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0245.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0246 = get_kernel('dispatch_kernel_0246') + + +def launch_dispatch_kernel_0246( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0246.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0247 = get_kernel('dispatch_kernel_0247') + + +def launch_dispatch_kernel_0247( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0247.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0248 = get_kernel('dispatch_kernel_0248') + + +def launch_dispatch_kernel_0248( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0248.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0249 = get_kernel('dispatch_kernel_0249') + + +def launch_dispatch_kernel_0249( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0249.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0250 = get_kernel('dispatch_kernel_0250') + + +def launch_dispatch_kernel_0250( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0250.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0251 = get_kernel('dispatch_kernel_0251') + + +def launch_dispatch_kernel_0251( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0251.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0252 = get_kernel('dispatch_kernel_0252') + + +def launch_dispatch_kernel_0252( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0252.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0253 = get_kernel('dispatch_kernel_0253') + + +def launch_dispatch_kernel_0253( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0253.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0254 = get_kernel('dispatch_kernel_0254') + + +def launch_dispatch_kernel_0254( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0254.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0255 = get_kernel('dispatch_kernel_0255') + + +def launch_dispatch_kernel_0255( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0255.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0256 = get_kernel('dispatch_kernel_0256') + + +def launch_dispatch_kernel_0256( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0256.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0257 = get_kernel('dispatch_kernel_0257') + + +def launch_dispatch_kernel_0257( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0257.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0258 = get_kernel('dispatch_kernel_0258') + + +def launch_dispatch_kernel_0258( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0258.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0259 = get_kernel('dispatch_kernel_0259') + + +def launch_dispatch_kernel_0259( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0259.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0260 = get_kernel('dispatch_kernel_0260') + + +def launch_dispatch_kernel_0260( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0260.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0261 = get_kernel('dispatch_kernel_0261') + + +def launch_dispatch_kernel_0261( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0261.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0262 = get_kernel('dispatch_kernel_0262') + + +def launch_dispatch_kernel_0262( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0262.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0263 = get_kernel('dispatch_kernel_0263') + + +def launch_dispatch_kernel_0263( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0263.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0264 = get_kernel('dispatch_kernel_0264') + + +def launch_dispatch_kernel_0264( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0264.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0265 = get_kernel('dispatch_kernel_0265') + + +def launch_dispatch_kernel_0265( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0265.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0266 = get_kernel('dispatch_kernel_0266') + + +def launch_dispatch_kernel_0266( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0266.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0267 = get_kernel('dispatch_kernel_0267') + + +def launch_dispatch_kernel_0267( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0267.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0268 = get_kernel('dispatch_kernel_0268') + + +def launch_dispatch_kernel_0268( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0268.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0269 = get_kernel('dispatch_kernel_0269') + + +def launch_dispatch_kernel_0269( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0269.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0270 = get_kernel('dispatch_kernel_0270') + + +def launch_dispatch_kernel_0270( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0270.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0271 = get_kernel('dispatch_kernel_0271') + + +def launch_dispatch_kernel_0271( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0271.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0272 = get_kernel('dispatch_kernel_0272') + + +def launch_dispatch_kernel_0272( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0272.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0273 = get_kernel('dispatch_kernel_0273') + + +def launch_dispatch_kernel_0273( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0273.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0274 = get_kernel('dispatch_kernel_0274') + + +def launch_dispatch_kernel_0274( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0274.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0275 = get_kernel('dispatch_kernel_0275') + + +def launch_dispatch_kernel_0275( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0275.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0276 = get_kernel('dispatch_kernel_0276') + + +def launch_dispatch_kernel_0276( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0276.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0277 = get_kernel('dispatch_kernel_0277') + + +def launch_dispatch_kernel_0277( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0277.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0278 = get_kernel('dispatch_kernel_0278') + + +def launch_dispatch_kernel_0278( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0278.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0279 = get_kernel('dispatch_kernel_0279') + + +def launch_dispatch_kernel_0279( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0279.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0280 = get_kernel('dispatch_kernel_0280') + + +def launch_dispatch_kernel_0280( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0280.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0281 = get_kernel('dispatch_kernel_0281') + + +def launch_dispatch_kernel_0281( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0281.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0282 = get_kernel('dispatch_kernel_0282') + + +def launch_dispatch_kernel_0282( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0282.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0283 = get_kernel('dispatch_kernel_0283') + + +def launch_dispatch_kernel_0283( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0283.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0284 = get_kernel('dispatch_kernel_0284') + + +def launch_dispatch_kernel_0284( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0284.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0285 = get_kernel('dispatch_kernel_0285') + + +def launch_dispatch_kernel_0285( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0285.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0286 = get_kernel('dispatch_kernel_0286') + + +def launch_dispatch_kernel_0286( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0286.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0287 = get_kernel('dispatch_kernel_0287') + + +def launch_dispatch_kernel_0287( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0287.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0288 = get_kernel('dispatch_kernel_0288') + + +def launch_dispatch_kernel_0288( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0288.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0289 = get_kernel('dispatch_kernel_0289') + + +def launch_dispatch_kernel_0289( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0289.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0290 = get_kernel('dispatch_kernel_0290') + + +def launch_dispatch_kernel_0290( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0290.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0291 = get_kernel('dispatch_kernel_0291') + + +def launch_dispatch_kernel_0291( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0291.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0292 = get_kernel('dispatch_kernel_0292') + + +def launch_dispatch_kernel_0292( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0292.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0293 = get_kernel('dispatch_kernel_0293') + + +def launch_dispatch_kernel_0293( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0293.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0294 = get_kernel('dispatch_kernel_0294') + + +def launch_dispatch_kernel_0294( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0294.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0295 = get_kernel('dispatch_kernel_0295') + + +def launch_dispatch_kernel_0295( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0295.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0296 = get_kernel('dispatch_kernel_0296') + + +def launch_dispatch_kernel_0296( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0296.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0297 = get_kernel('dispatch_kernel_0297') + + +def launch_dispatch_kernel_0297( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0297.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0298 = get_kernel('dispatch_kernel_0298') + + +def launch_dispatch_kernel_0298( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0298.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0299 = get_kernel('dispatch_kernel_0299') + + +def launch_dispatch_kernel_0299( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0299.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0300 = get_kernel('dispatch_kernel_0300') + + +def launch_dispatch_kernel_0300( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0300.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0301 = get_kernel('dispatch_kernel_0301') + + +def launch_dispatch_kernel_0301( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0301.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0302 = get_kernel('dispatch_kernel_0302') + + +def launch_dispatch_kernel_0302( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0302.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0303 = get_kernel('dispatch_kernel_0303') + + +def launch_dispatch_kernel_0303( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0303.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0304 = get_kernel('dispatch_kernel_0304') + + +def launch_dispatch_kernel_0304( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0304.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0305 = get_kernel('dispatch_kernel_0305') + + +def launch_dispatch_kernel_0305( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0305.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0306 = get_kernel('dispatch_kernel_0306') + + +def launch_dispatch_kernel_0306( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0306.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0307 = get_kernel('dispatch_kernel_0307') + + +def launch_dispatch_kernel_0307( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0307.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0308 = get_kernel('dispatch_kernel_0308') + + +def launch_dispatch_kernel_0308( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0308.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0309 = get_kernel('dispatch_kernel_0309') + + +def launch_dispatch_kernel_0309( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0309.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0310 = get_kernel('dispatch_kernel_0310') + + +def launch_dispatch_kernel_0310( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0310.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0311 = get_kernel('dispatch_kernel_0311') + + +def launch_dispatch_kernel_0311( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0311.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0312 = get_kernel('dispatch_kernel_0312') + + +def launch_dispatch_kernel_0312( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0312.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0313 = get_kernel('dispatch_kernel_0313') + + +def launch_dispatch_kernel_0313( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0313.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0314 = get_kernel('dispatch_kernel_0314') + + +def launch_dispatch_kernel_0314( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0314.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0315 = get_kernel('dispatch_kernel_0315') + + +def launch_dispatch_kernel_0315( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0315.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0316 = get_kernel('dispatch_kernel_0316') + + +def launch_dispatch_kernel_0316( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0316.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0317 = get_kernel('dispatch_kernel_0317') + + +def launch_dispatch_kernel_0317( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0317.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0318 = get_kernel('dispatch_kernel_0318') + + +def launch_dispatch_kernel_0318( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0318.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0319 = get_kernel('dispatch_kernel_0319') + + +def launch_dispatch_kernel_0319( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0319.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0320 = get_kernel('dispatch_kernel_0320') + + +def launch_dispatch_kernel_0320( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0320.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0321 = get_kernel('dispatch_kernel_0321') + + +def launch_dispatch_kernel_0321( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0321.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0322 = get_kernel('dispatch_kernel_0322') + + +def launch_dispatch_kernel_0322( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0322.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0323 = get_kernel('dispatch_kernel_0323') + + +def launch_dispatch_kernel_0323( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0323.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0324 = get_kernel('dispatch_kernel_0324') + + +def launch_dispatch_kernel_0324( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0324.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0325 = get_kernel('dispatch_kernel_0325') + + +def launch_dispatch_kernel_0325( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0325.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0326 = get_kernel('dispatch_kernel_0326') + + +def launch_dispatch_kernel_0326( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0326.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0327 = get_kernel('dispatch_kernel_0327') + + +def launch_dispatch_kernel_0327( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0327.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0328 = get_kernel('dispatch_kernel_0328') + + +def launch_dispatch_kernel_0328( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0328.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0329 = get_kernel('dispatch_kernel_0329') + + +def launch_dispatch_kernel_0329( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0329.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0330 = get_kernel('dispatch_kernel_0330') + + +def launch_dispatch_kernel_0330( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0330.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0331 = get_kernel('dispatch_kernel_0331') + + +def launch_dispatch_kernel_0331( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0331.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0332 = get_kernel('dispatch_kernel_0332') + + +def launch_dispatch_kernel_0332( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0332.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0333 = get_kernel('dispatch_kernel_0333') + + +def launch_dispatch_kernel_0333( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0333.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0334 = get_kernel('dispatch_kernel_0334') + + +def launch_dispatch_kernel_0334( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0334.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0335 = get_kernel('dispatch_kernel_0335') + + +def launch_dispatch_kernel_0335( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0335.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0336 = get_kernel('dispatch_kernel_0336') + + +def launch_dispatch_kernel_0336( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0336.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0337 = get_kernel('dispatch_kernel_0337') + + +def launch_dispatch_kernel_0337( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0337.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0338 = get_kernel('dispatch_kernel_0338') + + +def launch_dispatch_kernel_0338( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0338.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0339 = get_kernel('dispatch_kernel_0339') + + +def launch_dispatch_kernel_0339( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0339.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0340 = get_kernel('dispatch_kernel_0340') + + +def launch_dispatch_kernel_0340( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0340.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0341 = get_kernel('dispatch_kernel_0341') + + +def launch_dispatch_kernel_0341( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0341.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0342 = get_kernel('dispatch_kernel_0342') + + +def launch_dispatch_kernel_0342( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0342.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0343 = get_kernel('dispatch_kernel_0343') + + +def launch_dispatch_kernel_0343( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0343.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0344 = get_kernel('dispatch_kernel_0344') + + +def launch_dispatch_kernel_0344( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0344.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0345 = get_kernel('dispatch_kernel_0345') + + +def launch_dispatch_kernel_0345( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0345.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0346 = get_kernel('dispatch_kernel_0346') + + +def launch_dispatch_kernel_0346( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0346.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0347 = get_kernel('dispatch_kernel_0347') + + +def launch_dispatch_kernel_0347( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0347.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0348 = get_kernel('dispatch_kernel_0348') + + +def launch_dispatch_kernel_0348( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0348.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0349 = get_kernel('dispatch_kernel_0349') + + +def launch_dispatch_kernel_0349( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0349.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0350 = get_kernel('dispatch_kernel_0350') + + +def launch_dispatch_kernel_0350( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0350.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0351 = get_kernel('dispatch_kernel_0351') + + +def launch_dispatch_kernel_0351( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0351.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0352 = get_kernel('dispatch_kernel_0352') + + +def launch_dispatch_kernel_0352( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0352.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0353 = get_kernel('dispatch_kernel_0353') + + +def launch_dispatch_kernel_0353( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0353.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0354 = get_kernel('dispatch_kernel_0354') + + +def launch_dispatch_kernel_0354( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0354.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0355 = get_kernel('dispatch_kernel_0355') + + +def launch_dispatch_kernel_0355( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0355.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0356 = get_kernel('dispatch_kernel_0356') + + +def launch_dispatch_kernel_0356( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0356.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) + + +dispatch_kernel_0357 = get_kernel('dispatch_kernel_0357') + + +def launch_dispatch_kernel_0357( + *args, + grid: tuple[int, int, int], + block: tuple[int, int, int] | None = None, + shared_mem: int | None = None, + stream=None, + timeout_ms: float | None = None, + arch: str | None = None, + options: list[str] | None = None, +): + return dispatch_kernel_0357.launch( + *args, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=stream, + timeout_ms=timeout_ms, + arch=arch, + options=options, + ) diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/manifest.json b/cake_exports/kmeans/src/flashlib_cake_kmeans/manifest.json new file mode 100644 index 00000000..e1fd0bb4 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/manifest.json @@ -0,0 +1,28144 @@ +{ + "export_plan": { + "entrypoints": { + "correctness_test": "tests/test_correctness.py", + "performance_benchmark": "benchmarks/benchmark.py", + "python_interface": "src/flashlib_cake_kmeans/interface.py" + }, + "file": "plan.json", + "kernel_count": 358, + "metadata": { + "baseline": "same-invocation pinned 07cf Triton baseline measured with CUPTI", + "inventory": "complete production dispatcher ledger", + "launch_resolution": "init-once device runtime with lazy per-shape/per-stream pointer-rebindable direct sequences", + "registry_correctness_shapes": 124, + "registry_performance_shapes": 124, + "runtime_coverage_shapes": 104, + "workload": "kmeans_assignment" + }, + "name": "flashlib-kmeans-production", + "package_exports": { + "FlashKMeansAssignRuntime": ".interface:FlashKMeansAssignRuntime", + "PreparedFlashKMeansAssign": ".interface:PreparedFlashKMeansAssign", + "flash_kmeans_assign": ".interface:flash_kmeans_assign", + "flash_kmeans_assign_prepared": ".interface:flash_kmeans_assign_prepared", + "init": ".interface:init", + "prepare_flash_kmeans_assign": ".interface:prepare_flash_kmeans_assign" + }, + "sha256": "e5f7cbc5c8b6b666b8c9c546f06e82c94b8dff7049032332ea0b936a5df4eac1", + "tvm_ffi_exports": { + "flash_kmeans_assign": ".interface:flash_kmeans_assign" + }, + "validation": { + "benchmark_registry_baseline": { + "official": { + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_entrypoint": "benchmarks.flash_kmeans_triton_h200_raw_adapter:TritonH20007cfRawAdapter.compute", + "baseline_name": "triton_h200_07cf_raw_adapter_v1", + "baseline_timing_backend_field": "baseline_07cf_adapter_timing_backend", + "baseline_timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e", + "candidate_entrypoint": "flashlib_cake_kmeans.interface:FlashKMeansAssignRuntime.compute", + "candidate_timing_backend_field": "candidate_public_raw_timing_backend", + "candidate_timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion", + "role": "publication", + "speedup_denominator_metric": "candidate_public_raw_synchronized_e2e_ms", + "speedup_metric": "public_raw_e2e_speedup_vs_07cf_adapter", + "speedup_numerator_metric": "baseline_07cf_adapter_synchronized_e2e_ms", + "timing_backend": "cupti" + }, + "parity": { + "baseline_commit": "07cf2a27928aacf6790c950a265d8b8dc83c87cf", + "baseline_entrypoint": "benchmarks.flash_kmeans_triton_h200:euclid_assign_triton_h200", + "baseline_name": "triton_h200_07cf_precomputed", + "baseline_timing_backend_field": "baseline_07cf_precomputed_timing_backend", + "baseline_timing_boundary": "precomputed_norms_preallocated_output_pinned_07cf_assignment_gpu_span", + "candidate_entrypoint": "flashlib_cake_kmeans.interface:flash_kmeans_assign_prepared", + "candidate_timing_backend_field": "candidate_precomputed_timing_backend", + "candidate_timing_boundary": "precomputed_norms_preallocated_output_prepared_assignment_gpu_span", + "role": "diagnostic_only", + "speedup_denominator_metric": "candidate_precomputed_gpu_span_ms", + "speedup_metric": "precomputed_gpu_speedup_vs_07cf", + "speedup_numerator_metric": "baseline_07cf_precomputed_gpu_span_ms", + "timing_backend": "cupti" + }, + "registry_candidate_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shared_preprocess": { + "implementation_entrypoint": "flashlib_cake_kmeans._row_norm:PreparedBF16PairRowNorm", + "result_source_sha256_field": "preprocess_source_sha256", + "source_sha256": "aa67813cf1cc39b8ae96970a737e926f7b3a65dac63dbeb6362f2dacf066e26e" + } + }, + "benchmark_registry_baseline_key": "triton_h200_07cf_dual_lane_v1", + "benchmark_registry_baseline_sha256": "bdfd30338aa614f09817af1498c1115a0b1732a3340a1acbc172d0e4dd4674c4", + "benchmark_registry_candidate_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "benchmark_registry_key": "flash_kmeans_assign_dispatch_full_contract", + "benchmark_registry_shape_labels_sha256": "172741a45af34837224b4005e0b6ba45637a07c7a4cabc88ed85ec95670e6290", + "benchmark_registry_shape_ledger_sha256": "172741a45af34837224b4005e0b6ba45637a07c7a4cabc88ed85ec95670e6290", + "correctness_shape_count": 124, + "expected_routes": [ + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_1d49_d112_highk_tail_b1_n1408_k4096_d112" + }, + { + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_1d49_d128_forced_fallback_b5_n7296_k512_d128" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_1d49_d224_highn_overlap_b5_n5632_k768_d224" + }, + { + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_1d49_d288_low_n_heldout_b3_n1152_k2048_d288" + }, + { + "selected_route": "d352_exactd_splitk_c95c_v2", + "shape": "adjacent_1d49_d352_request_highk_b2_n1024_k8192_d352" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_1d49_d416_random_b4_n3840_k512_d416" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_1d49_d480_smallk_boundary_b2_n2816_k256_d480" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_1d49_d48_lowk_boundary_b4_n3968_k256_d48" + }, + { + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "shape": "adjacent_3328_d112_highk_low_n_b2_n512_k8192_d112" + }, + { + "selected_route": "paired_large_v15", + "shape": "adjacent_3328_d128_exact_guard_k_neighbor_b8_n8192_k512_d128" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_3328_d224_tail_div_b3_n3840_k512_d224" + }, + { + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_3328_d288_guard_overlap_b4_n8192_k256_d288" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_3328_d352_random_legal_b3_n2048_k768_d352" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_3328_d416_highk_low_n_b1_n384_k8192_d416" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_3328_d480_min_boundary_b1_n128_k256_d480" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_3328_d48_small_boundary_b1_n256_k256_d48" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_5600_d112_odd_tile_tail_b5_n2176_k512_d112" + }, + { + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_5600_d128_fallback_neighbor_b8_n8320_k256_d128" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_5600_d224_k512_overlap_b3_n3072_k512_d224" + }, + { + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_5600_d288_mid_highk_b2_n1024_k2048_d288" + }, + { + "selected_route": "d352_exactd_splitk_c95c_v2", + "shape": "adjacent_5600_d352_splitk_edge_b2_n768_k4096_d352" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_5600_d416_smallk_boundary_b4_n2048_k256_d416" + }, + { + "selected_route": "d480_splitk_k1024_eac2_v1", + "shape": "adjacent_5600_d480_random_b1_n1536_k1024_d480" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_5600_d48_low_n_boundary_b3_n1536_k256_d48" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_68cf_d112_tail_b5_n2944_k512_d112" + }, + { + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_68cf_d128_forced_fallback_b7_n6016_k512_d128" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_68cf_d224_overlap_b3_n5120_k1024_d224" + }, + { + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_68cf_d288_boundary_b2_n1920_k512_d288" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_68cf_d352_tail_b3_n2816_k768_d352" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_68cf_d416_overlap_b2_n2304_k1024_d416" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_68cf_d480_boundary_b4_n1664_k512_d480" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_68cf_d48_boundary_b4_n2304_k512_d48" + }, + { + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "shape": "adjacent_8f09_d112_small_highk_b1_n384_k4096_d112" + }, + { + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_8f09_d128_fallback_neighbor_b4_n4480_k512_d128" + }, + { + "selected_route": "d224_tmem_abi_repair_d17c_v4", + "shape": "adjacent_8f09_d224_low_n_boundary_b4_n1536_k256_d224" + }, + { + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_8f09_d288_k768_overlap_b1_n2560_k768_d288" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_8f09_d352_smallk_large_n_b4_n4096_k256_d352" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_8f09_d416_midk_tail_b3_n3456_k768_d416" + }, + { + "selected_route": "d480_splitk_k1024_eac2_v1", + "shape": "adjacent_8f09_d480_highk_low_n_b2_n640_k4096_d480" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_8f09_d48_medium_tail_b2_n1792_k512_d48" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_9ca0_d112_tail_div_b3_n3840_k768_d112" + }, + { + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_9ca0_d128_exact_guard_neighbor_b8_n8064_k256_d128" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_9ca0_d224_guard_overlap_b2_n4096_k256_d224" + }, + { + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_9ca0_d288_highk_low_n_b1_n384_k4096_d288" + }, + { + "selected_route": "d352_exactd_splitk_c95c_v2", + "shape": "adjacent_9ca0_d352_splitk_boundary_b1_n512_k8192_d352" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_9ca0_d416_splitk_min_b1_n1024_k4096_d416" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_9ca0_d480_splitd_large_n_b2_n4096_k512_d480" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_9ca0_d48_guard_boundary_b6_n12288_k512_d48" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_a2f8_d112_lowk_tail_b4_n3456_k256_d112" + }, + { + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_a2f8_d128_odd_fallback_b6_n8576_k512_d128" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_a2f8_d224_highn_overlap_b2_n6144_k768_d224" + }, + { + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_a2f8_d288_splitk_probe_b1_n640_k4096_d288" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_a2f8_d352_smallk_boundary_b4_n1024_k256_d352" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_a2f8_d416_random_b2_n2560_k768_d416" + }, + { + "selected_route": "d480_splitk_k1024_eac2_v1", + "shape": "adjacent_a2f8_d480_highk_request_b1_n896_k4096_d480" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_a2f8_d48_k1024_boundary_b2_n768_k1024_d48" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_c44f_d112_odd_tail_b2_n3200_k768_d112" + }, + { + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_c44f_d128_forced_fallback_b6_n6272_k512_d128" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_c44f_d224_overlap_b4_n4480_k512_d224" + }, + { + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_c44f_d288_highk_heldout_b2_n768_k4096_d288" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_c44f_d352_random_b1_n3328_k768_d352" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_c44f_d416_highk_request_b3_n512_k8192_d416" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_c44f_d480_boundary_b5_n2048_k512_d480" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_c44f_d48_midk_boundary_b3_n2688_k768_d48" + }, + { + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "shape": "adjacent_d9d5_d112_highk_request_b3_n768_k8192_d112" + }, + { + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_d9d5_d128_guard_miss_fallback_b5_n6016_k1024_d128" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_d9d5_d224_random_midk_b2_n2944_k1280_d224" + }, + { + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_d9d5_d288_heldout_low_n_b1_n896_k4096_d288" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_d9d5_d352_random_b5_n2304_k768_d352" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_d9d5_d416_highk_request_b2_n640_k8192_d416" + }, + { + "selected_route": "d480_splitk_k1024_eac2_v1", + "shape": "adjacent_d9d5_d480_heldout_highk_b3_n1024_k4096_d480" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_d9d5_d48_highk_heldout_b2_n640_k4096_d48" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_ef00_d112_odd_tail_b4_n3712_k1024_d112" + }, + { + "selected_route": "aligned_weave_v10_fallback", + "shape": "adjacent_ef00_d128_forced_fallback_b4_n7552_k512_d128" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_ef00_d224_midn_overlap_b2_n5376_k512_d224" + }, + { + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "adjacent_ef00_d288_highk_heldout_b1_n1664_k4096_d288" + }, + { + "selected_route": "d352_exactd_splitk_c95c_v2", + "shape": "adjacent_ef00_d352_request_low_n_b3_n896_k8192_d352" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_ef00_d416_random_midk_b1_n2176_k768_d416" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_ef00_d480_lowk_boundary_b3_n3200_k256_d480" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "adjacent_ef00_d48_midn_boundary_b3_n3456_k512_d48" + }, + { + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "shape": "post_d895_d112_b1_n256_k256_d112" + }, + { + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "shape": "post_d895_d112_b1_n512_k8192_d112" + }, + { + "selected_route": "d112_f826_c829_full_bucket_weave_evolve_flash_kmeans_assign_a262_v1", + "shape": "post_d895_d112_b2_n2048_k512_d112" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d112_b4_n8192_k1024_d112" + }, + { + "selected_route": "aligned_weave_v10_fallback", + "shape": "post_d895_d128_fallback_b3_n1920_k256_d128" + }, + { + "selected_route": "aligned_weave_v10_fallback", + "shape": "post_d895_d128_fallback_b5_n2176_k512_d128" + }, + { + "selected_route": "aligned_weave_v10_fallback", + "shape": "post_d895_d128_fallback_b7_n2432_k1024_d128" + }, + { + "selected_route": "paired_large_v15", + "shape": "post_d895_d128_paired_b2_n262144_k256_d128" + }, + { + "selected_route": "d128_even_near_floor_v10_repair", + "shape": "post_d895_d128_paired_b8_n8192_k256_d128" + }, + { + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d144_b1_n256_k256_d144" + }, + { + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d144_b1_n512_k8192_d144" + }, + { + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d144_b2_n2048_k1024_d144" + }, + { + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d144_b4_n8192_k1024_d144" + }, + { + "selected_route": "microdim_hybrid_9c0d_v1", + "shape": "post_d895_d16_b4_n32768_k1024_d16" + }, + { + "selected_route": "microdim_pipeline4_08f9_v4", + "shape": "post_d895_d16_b8_n65536_k512_d16" + }, + { + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d176_b1_n256_k256_d176" + }, + { + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d176_b1_n512_k8192_d176" + }, + { + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d176_b2_n2048_k1024_d176" + }, + { + "selected_route": "d144_d160_d176_pad192_tail_repair_f9b2_v1", + "shape": "post_d895_d176_b4_n8192_k1024_d176" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d224_b1_n256_k256_d224" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d224_b1_n512_k8192_d224" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d224_b2_n2048_k1024_d224" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d224_b4_n8192_k1024_d224" + }, + { + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "post_d895_d288_b1_n256_k256_d288" + }, + { + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "post_d895_d288_b1_n512_k8192_d288" + }, + { + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "post_d895_d288_b2_n2048_k1024_d288" + }, + { + "selected_route": "d288_parent_splitk_hybrid_20260629_v1", + "shape": "post_d895_d288_b4_n8192_k1024_d288" + }, + { + "selected_route": "microdim_hybrid_9c0d_v1", + "shape": "post_d895_d32_b4_n32768_k1024_d32" + }, + { + "selected_route": "microdim_pipeline4_08f9_v4", + "shape": "post_d895_d32_b8_n65536_k512_d32" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d352_b1_n256_k256_d352" + }, + { + "selected_route": "d352_exactd_splitk_c95c_v2", + "shape": "post_d895_d352_b1_n512_k8192_d352" + }, + { + "selected_route": "d352_exactd_splitk_c95c_v2", + "shape": "post_d895_d352_b2_n2048_k1024_d352" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d352_b4_n8192_k1024_d352" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d416_b1_n256_k256_d416" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d416_b1_n512_k8192_d416" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d416_b2_n2048_k1024_d416" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d416_b4_n8192_k1024_d416" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d480_b1_n256_k256_d480" + }, + { + "selected_route": "d480_splitk_k1024_eac2_v1", + "shape": "post_d895_d480_b1_n512_k8192_d480" + }, + { + "selected_route": "d480_splitk_k1024_eac2_v1", + "shape": "post_d895_d480_b2_n2048_k1024_d480" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d480_b4_n8192_k1024_d480" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d48_b1_n512_k8192_d48" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d48_b2_n2048_k512_d48" + }, + { + "selected_route": "gap_pad_to_supported_seed_v1", + "shape": "post_d895_d48_b4_n8192_k1024_d48" + } + ], + "expected_shape_count": 124, + "minimum_speedup": 1.0, + "performance_floor_scope": "all_benchmarked_shapes", + "performance_floor_shape_labels": [ + "adjacent_1d49_d112_highk_tail_b1_n1408_k4096_d112", + "adjacent_1d49_d128_forced_fallback_b5_n7296_k512_d128", + "adjacent_1d49_d224_highn_overlap_b5_n5632_k768_d224", + "adjacent_1d49_d288_low_n_heldout_b3_n1152_k2048_d288", + "adjacent_1d49_d352_request_highk_b2_n1024_k8192_d352", + "adjacent_1d49_d416_random_b4_n3840_k512_d416", + "adjacent_1d49_d480_smallk_boundary_b2_n2816_k256_d480", + "adjacent_1d49_d48_lowk_boundary_b4_n3968_k256_d48", + "adjacent_3328_d112_highk_low_n_b2_n512_k8192_d112", + "adjacent_3328_d128_exact_guard_k_neighbor_b8_n8192_k512_d128", + "adjacent_3328_d224_tail_div_b3_n3840_k512_d224", + "adjacent_3328_d288_guard_overlap_b4_n8192_k256_d288", + "adjacent_3328_d352_random_legal_b3_n2048_k768_d352", + "adjacent_3328_d416_highk_low_n_b1_n384_k8192_d416", + "adjacent_3328_d480_min_boundary_b1_n128_k256_d480", + "adjacent_3328_d48_small_boundary_b1_n256_k256_d48", + "adjacent_5600_d112_odd_tile_tail_b5_n2176_k512_d112", + "adjacent_5600_d128_fallback_neighbor_b8_n8320_k256_d128", + "adjacent_5600_d224_k512_overlap_b3_n3072_k512_d224", + "adjacent_5600_d288_mid_highk_b2_n1024_k2048_d288", + "adjacent_5600_d352_splitk_edge_b2_n768_k4096_d352", + "adjacent_5600_d416_smallk_boundary_b4_n2048_k256_d416", + "adjacent_5600_d480_random_b1_n1536_k1024_d480", + "adjacent_5600_d48_low_n_boundary_b3_n1536_k256_d48", + "adjacent_68cf_d112_tail_b5_n2944_k512_d112", + "adjacent_68cf_d128_forced_fallback_b7_n6016_k512_d128", + "adjacent_68cf_d224_overlap_b3_n5120_k1024_d224", + "adjacent_68cf_d288_boundary_b2_n1920_k512_d288", + "adjacent_68cf_d352_tail_b3_n2816_k768_d352", + "adjacent_68cf_d416_overlap_b2_n2304_k1024_d416", + "adjacent_68cf_d480_boundary_b4_n1664_k512_d480", + "adjacent_68cf_d48_boundary_b4_n2304_k512_d48", + "adjacent_8f09_d112_small_highk_b1_n384_k4096_d112", + "adjacent_8f09_d128_fallback_neighbor_b4_n4480_k512_d128", + "adjacent_8f09_d224_low_n_boundary_b4_n1536_k256_d224", + "adjacent_8f09_d288_k768_overlap_b1_n2560_k768_d288", + "adjacent_8f09_d352_smallk_large_n_b4_n4096_k256_d352", + "adjacent_8f09_d416_midk_tail_b3_n3456_k768_d416", + "adjacent_8f09_d480_highk_low_n_b2_n640_k4096_d480", + "adjacent_8f09_d48_medium_tail_b2_n1792_k512_d48", + "adjacent_9ca0_d112_tail_div_b3_n3840_k768_d112", + "adjacent_9ca0_d128_exact_guard_neighbor_b8_n8064_k256_d128", + "adjacent_9ca0_d224_guard_overlap_b2_n4096_k256_d224", + "adjacent_9ca0_d288_highk_low_n_b1_n384_k4096_d288", + "adjacent_9ca0_d352_splitk_boundary_b1_n512_k8192_d352", + "adjacent_9ca0_d416_splitk_min_b1_n1024_k4096_d416", + "adjacent_9ca0_d480_splitd_large_n_b2_n4096_k512_d480", + "adjacent_9ca0_d48_guard_boundary_b6_n12288_k512_d48", + "adjacent_a2f8_d112_lowk_tail_b4_n3456_k256_d112", + "adjacent_a2f8_d128_odd_fallback_b6_n8576_k512_d128", + "adjacent_a2f8_d224_highn_overlap_b2_n6144_k768_d224", + "adjacent_a2f8_d288_splitk_probe_b1_n640_k4096_d288", + "adjacent_a2f8_d352_smallk_boundary_b4_n1024_k256_d352", + "adjacent_a2f8_d416_random_b2_n2560_k768_d416", + "adjacent_a2f8_d480_highk_request_b1_n896_k4096_d480", + "adjacent_a2f8_d48_k1024_boundary_b2_n768_k1024_d48", + "adjacent_c44f_d112_odd_tail_b2_n3200_k768_d112", + "adjacent_c44f_d128_forced_fallback_b6_n6272_k512_d128", + "adjacent_c44f_d224_overlap_b4_n4480_k512_d224", + "adjacent_c44f_d288_highk_heldout_b2_n768_k4096_d288", + "adjacent_c44f_d352_random_b1_n3328_k768_d352", + "adjacent_c44f_d416_highk_request_b3_n512_k8192_d416", + "adjacent_c44f_d480_boundary_b5_n2048_k512_d480", + "adjacent_c44f_d48_midk_boundary_b3_n2688_k768_d48", + "adjacent_d9d5_d112_highk_request_b3_n768_k8192_d112", + "adjacent_d9d5_d128_guard_miss_fallback_b5_n6016_k1024_d128", + "adjacent_d9d5_d224_random_midk_b2_n2944_k1280_d224", + "adjacent_d9d5_d288_heldout_low_n_b1_n896_k4096_d288", + "adjacent_d9d5_d352_random_b5_n2304_k768_d352", + "adjacent_d9d5_d416_highk_request_b2_n640_k8192_d416", + "adjacent_d9d5_d480_heldout_highk_b3_n1024_k4096_d480", + "adjacent_d9d5_d48_highk_heldout_b2_n640_k4096_d48", + "adjacent_ef00_d112_odd_tail_b4_n3712_k1024_d112", + "adjacent_ef00_d128_forced_fallback_b4_n7552_k512_d128", + "adjacent_ef00_d224_midn_overlap_b2_n5376_k512_d224", + "adjacent_ef00_d288_highk_heldout_b1_n1664_k4096_d288", + "adjacent_ef00_d352_request_low_n_b3_n896_k8192_d352", + "adjacent_ef00_d416_random_midk_b1_n2176_k768_d416", + "adjacent_ef00_d480_lowk_boundary_b3_n3200_k256_d480", + "adjacent_ef00_d48_midn_boundary_b3_n3456_k512_d48", + "post_d895_d112_b1_n256_k256_d112", + "post_d895_d112_b1_n512_k8192_d112", + "post_d895_d112_b2_n2048_k512_d112", + "post_d895_d112_b4_n8192_k1024_d112", + "post_d895_d128_fallback_b3_n1920_k256_d128", + "post_d895_d128_fallback_b5_n2176_k512_d128", + "post_d895_d128_fallback_b7_n2432_k1024_d128", + "post_d895_d128_paired_b2_n262144_k256_d128", + "post_d895_d128_paired_b8_n8192_k256_d128", + "post_d895_d144_b1_n256_k256_d144", + "post_d895_d144_b1_n512_k8192_d144", + "post_d895_d144_b2_n2048_k1024_d144", + "post_d895_d144_b4_n8192_k1024_d144", + "post_d895_d16_b4_n32768_k1024_d16", + "post_d895_d16_b8_n65536_k512_d16", + "post_d895_d176_b1_n256_k256_d176", + "post_d895_d176_b1_n512_k8192_d176", + "post_d895_d176_b2_n2048_k1024_d176", + "post_d895_d176_b4_n8192_k1024_d176", + "post_d895_d224_b1_n256_k256_d224", + "post_d895_d224_b1_n512_k8192_d224", + "post_d895_d224_b2_n2048_k1024_d224", + "post_d895_d224_b4_n8192_k1024_d224", + "post_d895_d288_b1_n256_k256_d288", + "post_d895_d288_b1_n512_k8192_d288", + "post_d895_d288_b2_n2048_k1024_d288", + "post_d895_d288_b4_n8192_k1024_d288", + "post_d895_d32_b4_n32768_k1024_d32", + "post_d895_d32_b8_n65536_k512_d32", + "post_d895_d352_b1_n256_k256_d352", + "post_d895_d352_b1_n512_k8192_d352", + "post_d895_d352_b2_n2048_k1024_d352", + "post_d895_d352_b4_n8192_k1024_d352", + "post_d895_d416_b1_n256_k256_d416", + "post_d895_d416_b1_n512_k8192_d416", + "post_d895_d416_b2_n2048_k1024_d416", + "post_d895_d416_b4_n8192_k1024_d416", + "post_d895_d480_b1_n256_k256_d480", + "post_d895_d480_b1_n512_k8192_d480", + "post_d895_d480_b2_n2048_k1024_d480", + "post_d895_d480_b4_n8192_k1024_d480", + "post_d895_d48_b1_n512_k8192_d48", + "post_d895_d48_b2_n2048_k512_d48", + "post_d895_d48_b4_n8192_k1024_d48" + ], + "performance_speedup_denominator_metric": "candidate_public_raw_synchronized_e2e_ms", + "performance_speedup_metric": "public_raw_e2e_speedup_vs_07cf_adapter", + "performance_speedup_numerator_metric": "baseline_07cf_adapter_synchronized_e2e_ms", + "runtime_lifecycle": { + "amortization_call_counts": [ + 1, + 10, + 100, + 1000 + ], + "baseline_api": "triton_h200_07cf_raw_adapter_v1.compute", + "baseline_code_cache_state": "process_order_dependent", + "baseline_first_cache_state": "shape_slot_miss", + "baseline_has_explicit_init": true, + "baseline_hot_cache_state": "fresh_pointer_shape_slot_hit", + "baseline_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "baseline_result_field": "baseline_runtime_lifecycle", + "baseline_timing_boundary": "raw_inputs_default_output_shared_fused_pair_norm_plus_pinned_07cf_assign_synchronized_e2e", + "cache_policy": "synchronize_and_clear_after_each_completed_shape", + "candidate_api": "flashlib_cake_kmeans.init(...).compute", + "candidate_code_cache_state": "process_order_dependent", + "candidate_first_cache_state": "shape_slot_miss", + "candidate_hot_cache_state": "fresh_pointer_shape_slot_hit", + "candidate_output_policy": "default_output_allocated_before_preprocessing_inside_timing", + "candidate_result_field": "candidate_runtime_lifecycle", + "candidate_timing_boundary": "raw_inputs_default_output_internal_required_norms_cached_dispatch_and_synchronized_completion", + "cold_order_policy": "deterministic_balanced_per_publication_contract_portfolio", + "comparison_result_field": "runtime_lifecycle_comparison", + "init_composition": "runtime_init_plus_shared_preprocess_support_each_standalone_lane", + "init_order_policy": "alternate_candidate_baseline_first_by_validation_shard_parity_then_shared_support", + "init_scope": "runtime_init_plus_standalone_shared_preprocess_support_once_per_validation_shard_process_device_operator", + "resident_multi_shape_cache_benchmarked": false, + "schema": "loom-public-runtime-lifecycle-v1" + }, + "semantic_entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "shape_test_marker": "export_validation_shape" + } + }, + "format_version": 1, + "generated_at_utc": "2026-07-09T04:26:21.680464+00:00", + "generator": "loom.export.kernel_repo", + "kernels": [ + { + "attr": "dispatch_kernel_0000", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_lowdim_pack_e50c_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0000", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0000.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_lowdim_pack_e50c_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0001", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_lowdim_e50c_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0001", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0001.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_lowdim_e50c_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0002", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0002", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0002.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0003", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0003", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0003.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0004", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0004", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0004.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0005", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0005", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0005.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0006", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0006", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0006.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0007", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0007", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0007.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0008", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0008", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0008.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0009", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0009", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0009.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0010", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0010", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0010.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0011", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitk_partial_blockn64_g2r4_b5a6_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0011", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_slices" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0011.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitk_partial_blockn64_g2r4_b5a6_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0012", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitk_reduce_blockn64_g2r4_b5a6_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0012", + "parameters": [ + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_slices" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0012.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitk_reduce_blockn64_g2r4_b5a6_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0013", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_microdim_pack_6cd2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0013", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0013.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_microdim_pack_6cd2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0014", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_microdim_6cd2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0014", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0014.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_microdim_6cd2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0015", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0015", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0015.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0016", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0016", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0016.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0017", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0017", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0017.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0018", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0018", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0018.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0019", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_microdim_pack_6cd2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0019", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0019.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_microdim_pack_6cd2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0020", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_microdim_6cd2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0020", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0020.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_microdim_6cd2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0021", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_microdim_pack_6cd2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0021", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0021.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_microdim_pack_6cd2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0022", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_microdim_6cd2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0022", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0022.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_microdim_6cd2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0023", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_microdim_pack_6cd2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0023", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0023.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_microdim_pack_6cd2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0024", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_microdim_6cd2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0024", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0024.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_microdim_6cd2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0025", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0025", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0025.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0026", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0026", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0026.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0027", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0027", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0027.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0028", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0028", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0028.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0029", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0029", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0029.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0030", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0030", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0030.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0031", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0031", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0031.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0032", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0032", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0032.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0033", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_lowdim_pack_e50c_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0033", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0033.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_lowdim_pack_e50c_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0034", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_lowdim_e50c_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0034", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0034.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_lowdim_e50c_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0035", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_lowdim_pack_e50c_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0035", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0035.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_lowdim_pack_e50c_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0036", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_lowdim_e50c_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0036", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0036.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_lowdim_e50c_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0037", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_lowdim_pack_e50c_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0037", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0037.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_lowdim_pack_e50c_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0038", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_lowdim_e50c_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0038", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0038.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_lowdim_e50c_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0039", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_lowdim_pack_e50c_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0039", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0039.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_lowdim_pack_e50c_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0040", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_lowdim_e50c_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0040", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0040.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_lowdim_e50c_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0041", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_lowdim_pack_e50c_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0041", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0041.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_lowdim_pack_e50c_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0042", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_lowdim_e50c_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0042", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0042.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_lowdim_e50c_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0043", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_lowdim_pack_e50c_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0043", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0043.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_lowdim_pack_e50c_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0044", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_lowdim_e50c_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0044", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0044.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_lowdim_e50c_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0045", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0045", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0045.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0046", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0046", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0046.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0047", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0047", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0047.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0048", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0048", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0048.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0049", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0049", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0049.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0050", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0050", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0050.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0051", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0051", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0051.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0052", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0052", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0052.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0053", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0053", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0053.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0054", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0054", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0054.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0055", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0055", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0055.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0056", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0056", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0056.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0057", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0057", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0057.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0058", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0058", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0058.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0059", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitk_partial_blockn64_g1r4_streamdep_r63_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0059", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_slices" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0059.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitk_partial_blockn64_g1r4_streamdep_r63_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0060", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitk_reduce_blockn64_g1r4_streamdep_r63_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0060", + "parameters": [ + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_slices" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0060.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitk_reduce_blockn64_g1r4_streamdep_r63_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0061", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0061", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0061.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0062", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0062", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0062.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0063", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0063", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0063.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0064", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0064", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0064.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0065", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0065", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0065.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0066", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0066", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0066.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0067", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0067", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0067.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0068", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0068", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0068.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0069", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0069", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0069.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0070", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0070", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0070.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0071", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0071", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0071.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0072", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0072", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0072.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0073", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0073", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0073.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0074", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0074", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0074.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0075", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0075", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0075.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0076", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0076", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0076.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0077", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0077", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0077.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0078", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0078", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0078.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0079", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0079", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0079.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0080", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0080", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0080.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0081", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0081", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0081.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0082", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0082", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0082.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0083", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0083", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0083.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0084", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0084", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0084.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0085", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0085", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0085.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0086", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0086", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0086.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0087", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0087", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0087.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0088", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0088", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0088.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0089", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0089", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0089.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0090", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0090", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0090.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0091", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0091", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0091.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0092", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0092", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0092.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0093", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0093", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0093.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0094", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0094", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0094.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0095", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0095", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0095.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0096", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0096", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0096.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0097", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0097", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0097.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0098", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0098", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0098.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0099", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0099", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0099.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0100", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0100", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0100.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0101", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0101", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0101.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0102", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0102", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0102.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0103", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0103", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0103.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0104", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0104", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0104.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0105", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0105", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0105.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0106", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0106", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0106.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0107", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0107", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0107.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0108", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0108", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0108.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0109", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0109", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0109.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0110", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0110", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0110.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0111", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0111", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0111.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0112", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0112", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0112.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0113", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0113", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0113.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0114", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0114", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0114.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0115", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0115", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0115.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0116", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0116", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0116.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0117", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0117", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0117.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0118", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_paired_xreuse_dualtmem_producer_r47_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0118", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "unsigned long long*", + "kind": "pointer", + "name": "partial_keys" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_slices" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0118.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_paired_xreuse_dualtmem_producer_r47_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0119", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_paired_ownerreduce_r39_reduce1_unroll_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0119", + "parameters": [ + { + "ctype": "unsigned long long*", + "kind": "pointer", + "name": "partial_keys" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_slices" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0119.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_paired_ownerreduce_r39_reduce1_unroll_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0120", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0120", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0120.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0121", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_paired_packedpartial_producer_7b3c_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0121", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "unsigned long long*", + "kind": "pointer", + "name": "partial_keys" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_slices" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0121.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_paired_packedpartial_producer_7b3c_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0122", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_paired_packedpartial_reduce_r2_7b3c_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0122", + "parameters": [ + { + "ctype": "unsigned long long*", + "kind": "pointer", + "name": "partial_keys" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_slices" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0122.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_paired_packedpartial_reduce_r2_7b3c_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0123", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0123", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0123.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0124", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_microdim_pack_6cd2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0124", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0124.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_microdim_pack_6cd2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0125", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_microdim_6cd2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0125", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0125.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_microdim_6cd2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0126", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_microdim_d16_pipeline4_08f9_v4", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0126", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0126.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_microdim_d16_pipeline4_08f9_v4", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0127", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_microdim_pack_6cd2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0127", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0127.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_microdim_pack_6cd2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0128", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_microdim_6cd2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0128", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0128.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_microdim_6cd2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0129", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_microdim_raw_tma_08f9_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0129", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0129.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_microdim_raw_tma_08f9_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0130", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0130", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0130.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0131", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0131", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0131.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0132", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0132", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0132.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0133", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0133", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0133.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0134", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0134", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0134.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0135", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0135", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0135.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0136", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0136", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0136.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0137", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0137", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0137.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0138", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0138", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0138.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0139", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0139", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0139.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0140", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0140", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0140.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0141", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0141", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0141.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0142", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0142", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0142.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0143", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0143", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0143.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0144", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0144", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0144.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0145", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0145", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0145.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0146", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0146", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0146.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0147", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0147", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0147.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0148", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0148", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0148.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0149", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0149", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0149.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0150", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0150", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0150.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0151", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0151", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0151.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0152", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0152", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0152.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0153", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0153", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0153.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0154", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0154", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0154.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0155", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0155", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0155.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0156", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0156", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0156.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0157", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0157", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0157.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0158", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0158", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0158.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0159", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0159", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0159.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0160", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0160", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0160.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0161", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0161", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0161.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0162", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0162", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0162.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0163", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0163", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0163.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0164", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0164", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0164.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0165", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0165", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0165.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0166", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0166", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0166.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0167", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0167", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0167.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0168", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0168", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0168.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0169", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0169", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0169.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0170", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0170", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0170.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0171", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0171", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0171.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0172", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0172", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0172.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0173", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0173", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0173.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0174", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0174", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0174.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0175", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0175", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0175.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0176", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0176", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0176.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0177", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0177", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0177.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0178", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0178", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0178.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0179", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0179", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0179.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0180", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0180", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0180.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0181", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0181", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0181.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0182", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0182", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0182.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0183", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0183", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0183.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0184", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0184", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0184.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0185", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0185", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0185.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0186", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0186", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0186.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0187", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0187", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0187.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0188", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0188", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0188.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0189", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0189", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0189.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0190", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0190", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0190.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0191", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0191", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0191.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0192", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0192", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0192.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0193", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0193", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0193.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0194", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0194", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0194.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0195", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0195", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0195.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0196", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0196", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0196.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0197", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0197", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0197.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0198", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0198", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0198.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0199", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0199", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0199.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0200", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0200", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0200.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0201", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0201", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0201.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0202", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0202", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0202.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0203", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0203", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0203.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0204", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0204", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0204.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0205", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0205", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0205.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0206", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0206", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0206.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0207", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0207", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0207.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0208", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0208", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0208.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0209", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0209", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0209.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0210", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0210", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0210.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0211", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0211", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0211.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0212", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0212", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0212.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0213", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0213", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0213.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0214", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0214", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0214.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0215", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0215", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0215.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0216", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0216", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0216.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0217", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0217", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0217.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0218", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0218", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0218.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0219", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0219", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0219.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0220", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0220", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0220.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0221", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0221", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0221.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0222", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0222", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0222.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0223", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0223", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0223.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0224", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0224", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0224.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0225", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0225", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0225.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0226", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0226", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0226.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0227", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0227", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0227.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0228", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0228", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0228.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0229", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0229", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0229.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0230", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0230", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0230.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0231", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0231", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0231.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0232", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0232", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0232.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0233", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0233", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0233.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0234", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0234", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0234.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0235", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0235", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0235.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0236", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_exactd_a532_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0236", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0236.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_exactd_a532_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0237", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_splitk_cta_0438_v1_partial", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0237", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0237.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0238", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0238", + "parameters": [ + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0238.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0239", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_splitk_cta_0438_v1_partial", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0239", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0239.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0240", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0240", + "parameters": [ + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0240.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0241", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_splitk_cta_0438_v1_partial", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0241", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0241.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0242", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0242", + "parameters": [ + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0242.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0243", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_splitk_cta_0438_v1_partial", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0243", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0243.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0244", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0244", + "parameters": [ + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0244.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0245", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_splitk_cta_0438_v1_partial", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0245", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0245.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0246", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0246", + "parameters": [ + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0246.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0247", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_exactd_a532_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0247", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0247.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_exactd_a532_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0248", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_splitk_cta_0438_v1_partial", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0248", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0248.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0249", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0249", + "parameters": [ + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0249.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0250", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_exactd_a532_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0250", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0250.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_exactd_a532_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0251", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_exactd_a532_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0251", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0251.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_exactd_a532_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0252", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_exactd_a532_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0252", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0252.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_exactd_a532_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0253", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_exactd_a532_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0253", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0253.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_exactd_a532_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0254", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_exactd_a532_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0254", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0254.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_exactd_a532_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0255", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_exactd_a532_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0255", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0255.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_exactd_a532_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0256", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0256", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0256.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0257", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0257", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0257.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0258", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d480_splitk_partial_d32k256_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0258", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_slices" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0258.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d480_splitk_partial_d32k256_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0259", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d480_splitk_reduce_d32k256_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0259", + "parameters": [ + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_slices" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0259.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d480_splitk_reduce_d32k256_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0260", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0260", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0260.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0261", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0261", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0261.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0262", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0262", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0262.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0263", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0263", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0263.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0264", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0264", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0264.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0265", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0265", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0265.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0266", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0266", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0266.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0267", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0267", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0267.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0268", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0268", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0268.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0269", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0269", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0269.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0270", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0270", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0270.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0271", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0271", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0271.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0272", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0272", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0272.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0273", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0273", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0273.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0274", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0274", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0274.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0275", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0275", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0275.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0276", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0276", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0276.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0277", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitk_partial_8de8_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0277", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0277.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0278", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitk_reduce_8de8_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0278", + "parameters": [ + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0278.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0279", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0279", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0279.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0280", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitk_partial_8de8_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0280", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0280.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0281", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitk_reduce_8de8_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0281", + "parameters": [ + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0281.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0282", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0282", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0282.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0283", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitk_partial_8de8_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0283", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0283.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0284", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitk_reduce_8de8_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0284", + "parameters": [ + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0284.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0285", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0285", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0285.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0286", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0286", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0286.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0287", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0287", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0287.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0288", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitk_partial_8de8_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0288", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0288.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0289", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitk_reduce_8de8_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0289", + "parameters": [ + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0289.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0290", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0290", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0290.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0291", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0291", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0291.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0292", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0292", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0292.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0293", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0293", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0293.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0294", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0294", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0294.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0295", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0295", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0295.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0296", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0296", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0296.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0297", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitk_partial_8de8_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0297", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0297.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0298", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitk_reduce_8de8_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0298", + "parameters": [ + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0298.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0299", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0299", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0299.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0300", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0300", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0300.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0301", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0301", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0301.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0302", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0302", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0302.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0303", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0303", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0303.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0304", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0304", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0304.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0305", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0305", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0305.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0306", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0306", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0306.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0307", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0307", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0307.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0308", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0308", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0308.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0309", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0309", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0309.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0310", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0310", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0310.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0311", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0311", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0311.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0312", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0312", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0312.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0313", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0313", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0313.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0314", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0314", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0314.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0315", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0315", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0315.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0316", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0316", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0316.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0317", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0317", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0317.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0318", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0318", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0318.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0319", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0319", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0319.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0320", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0320", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0320.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0321", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0321", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0321.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0322", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0322", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0322.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0323", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitk_partial_blockn64_g2r4_b5a6_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0323", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_slices" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0323.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitk_partial_blockn64_g2r4_b5a6_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0324", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitk_reduce_blockn64_g2r4_b5a6_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0324", + "parameters": [ + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_slices" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0324.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitk_reduce_blockn64_g2r4_b5a6_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0325", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitk_partial_8de8_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0325", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0325.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitk_partial_8de8_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0326", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitk_reduce_8de8_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0326", + "parameters": [ + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0326.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitk_reduce_8de8_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0327", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitd_6fcf_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0327", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0327.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitd_6fcf_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0328", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitk_partial_blockn64_g1r4_streamdep_r63_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0328", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_slices" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0328.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitk_partial_blockn64_g1r4_streamdep_r63_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0329", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_highd_splitk_reduce_blockn64_g1r4_streamdep_r63_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0329", + "parameters": [ + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_slices" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0329.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_highd_splitk_reduce_blockn64_g1r4_streamdep_r63_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0330", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_microdim_d16_pipeline4_08f9_v4", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0330", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0330.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_microdim_d16_pipeline4_08f9_v4", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0331", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_microdim_raw_tma_08f9_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0331", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0331.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_microdim_raw_tma_08f9_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0332", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_microdim_direct_9c0d_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0332", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0332.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_microdim_direct_9c0d_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0333", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_microdim_pack_6cd2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0333", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0333.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_microdim_pack_6cd2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0334", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_microdim_6cd2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0334", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0334.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_microdim_6cd2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0335", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_lowdim_pack_e50c_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0335", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0335.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_lowdim_pack_e50c_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0336", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_microdim_6cd2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0336", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0336.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_microdim_6cd2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0337", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_lowdim_e50c_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0337", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0337.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_lowdim_e50c_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0338", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_gap_pad_pack_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0338", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "D_PAD" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0338.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_gap_pad_pack_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0339", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0339", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0339.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d64_direct_1p2gap_9f2a_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0340", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v15", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0340", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0340.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v15", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0341", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_v10", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0341", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0341.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_v10", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0342", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0342", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0342.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d256_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0343", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0343", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0343.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_splitd_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0344", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0344", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0344.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d192_single_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0345", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d480_splitk_partial_d32k256_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0345", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_slices" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0345.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d480_splitk_partial_d32k256_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0346", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d480_splitk_reduce_d32k256_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0346", + "parameters": [ + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_slices" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0346.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d480_splitk_reduce_d32k256_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0347", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d416_exactd_splitd_a4a579d1_v2", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0347", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0347.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d416_exactd_splitd_a4a579d1_v2", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0348", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_splitk_cta_0438_v1_partial", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0348", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0348.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_partial", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0349", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0349", + "parameters": [ + { + "ctype": "float*", + "kind": "pointer", + "name": "partial_scores" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "partial_indices" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0349.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_splitk_cta_0438_v1_reduce", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0350", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d288_exactd_a532_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0350", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0350.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d288_exactd_a532_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0351", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0351", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "x_sq" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0351.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d224_tmem_abi_repair_d17c_v4", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0352", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0352", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0352.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d160_pad192_pack_f9b2_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0353", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d112_k1024_owner_local_warp_mma_distinct_workfeed_c829_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0353", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0353.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d112_k1024_owner_local_warp_mma_distinct_workfeed_c829_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0354", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d112_k512_two_owner_peer_only_mma_f826_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0354", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0354.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d112_k512_two_owner_peer_only_mma_f826_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0355", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d112_shared_point_dual_issuer_tcgen05_6e1e_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0355", + "parameters": [ + { + "ctype": "const void*", + "kind": "pointer", + "name": "x_tmap" + }, + { + "ctype": "const void*", + "kind": "pointer", + "name": "c_tmap" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + }, + { + "ctype": "int", + "kind": "value", + "name": "K_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0355.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d112_shared_point_dual_issuer_tcgen05_6e1e_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0356", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0356", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "float*", + "kind": "pointer", + "name": "c_sq" + }, + { + "ctype": "int*", + "kind": "pointer", + "name": "out" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "num_n_tiles" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0356.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_d112_consumer_scoped_direct_mma_view_handoff_e5e1_v1", + "threads": 256 + }, + { + "attr": "dispatch_kernel_0357", + "cluster_dims": [ + 1, + 1, + 1 + ], + "compile_options": [ + "--use_fast_math" + ], + "ir_name": "flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1", + "launch_mode": "standard", + "module": "loom.examples.weave.flash_kmeans_assign_dispatcher", + "name": "dispatch_kernel_0357", + "parameters": [ + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "centroids" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "x_pad" + }, + { + "ctype": "__nv_bfloat16*", + "kind": "pointer", + "name": "c_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "B" + }, + { + "ctype": "int", + "kind": "value", + "name": "N" + }, + { + "ctype": "int", + "kind": "value", + "name": "D" + }, + { + "ctype": "int", + "kind": "value", + "name": "K" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_x_pad" + }, + { + "ctype": "int", + "kind": "value", + "name": "total_c_pad" + } + ], + "shared_mem_bytes": 0, + "source": "cuda/dispatch_kernel_0357.cu", + "specializations": {}, + "symbol": "kernel_flash_kmeans_assign_cleanroom_tcgen05_d160_pack_padded_b23d_v1", + "threads": 256 + } + ], + "package": "flashlib_cake_kmeans", + "semantic_dispatcher": { + "entrypoint": "loom.examples.weave.flash_kmeans_assign_dispatcher:launch_for_eval", + "exact_route_key_count": 0, + "kernel_count": 358, + "module_count": 52, + "route_decision_count": 0 + }, + "source_commit": "42070e96d0734cb580854baef60f17625ba33bb5", + "source_repo": "ssh://git@gitlab-master.nvidia.com:12051/cake/cake.git", + "tvm_ffi": { + "adapter": "tvm_ffi.py", + "default_namespace": "flashlib_cake_kmeans", + "low_level_functions": [ + "flashlib_cake_kmeans.launch_dispatch_kernel_0000", + "flashlib_cake_kmeans.launch_dispatch_kernel_0001", + "flashlib_cake_kmeans.launch_dispatch_kernel_0002", + "flashlib_cake_kmeans.launch_dispatch_kernel_0003", + "flashlib_cake_kmeans.launch_dispatch_kernel_0004", + "flashlib_cake_kmeans.launch_dispatch_kernel_0005", + "flashlib_cake_kmeans.launch_dispatch_kernel_0006", + "flashlib_cake_kmeans.launch_dispatch_kernel_0007", + "flashlib_cake_kmeans.launch_dispatch_kernel_0008", + "flashlib_cake_kmeans.launch_dispatch_kernel_0009", + "flashlib_cake_kmeans.launch_dispatch_kernel_0010", + "flashlib_cake_kmeans.launch_dispatch_kernel_0011", + "flashlib_cake_kmeans.launch_dispatch_kernel_0012", + "flashlib_cake_kmeans.launch_dispatch_kernel_0013", + "flashlib_cake_kmeans.launch_dispatch_kernel_0014", + "flashlib_cake_kmeans.launch_dispatch_kernel_0015", + "flashlib_cake_kmeans.launch_dispatch_kernel_0016", + "flashlib_cake_kmeans.launch_dispatch_kernel_0017", + "flashlib_cake_kmeans.launch_dispatch_kernel_0018", + "flashlib_cake_kmeans.launch_dispatch_kernel_0019", + "flashlib_cake_kmeans.launch_dispatch_kernel_0020", + "flashlib_cake_kmeans.launch_dispatch_kernel_0021", + "flashlib_cake_kmeans.launch_dispatch_kernel_0022", + "flashlib_cake_kmeans.launch_dispatch_kernel_0023", + "flashlib_cake_kmeans.launch_dispatch_kernel_0024", + "flashlib_cake_kmeans.launch_dispatch_kernel_0025", + "flashlib_cake_kmeans.launch_dispatch_kernel_0026", + "flashlib_cake_kmeans.launch_dispatch_kernel_0027", + "flashlib_cake_kmeans.launch_dispatch_kernel_0028", + "flashlib_cake_kmeans.launch_dispatch_kernel_0029", + "flashlib_cake_kmeans.launch_dispatch_kernel_0030", + "flashlib_cake_kmeans.launch_dispatch_kernel_0031", + "flashlib_cake_kmeans.launch_dispatch_kernel_0032", + "flashlib_cake_kmeans.launch_dispatch_kernel_0033", + "flashlib_cake_kmeans.launch_dispatch_kernel_0034", + "flashlib_cake_kmeans.launch_dispatch_kernel_0035", + "flashlib_cake_kmeans.launch_dispatch_kernel_0036", + "flashlib_cake_kmeans.launch_dispatch_kernel_0037", + "flashlib_cake_kmeans.launch_dispatch_kernel_0038", + "flashlib_cake_kmeans.launch_dispatch_kernel_0039", + "flashlib_cake_kmeans.launch_dispatch_kernel_0040", + "flashlib_cake_kmeans.launch_dispatch_kernel_0041", + "flashlib_cake_kmeans.launch_dispatch_kernel_0042", + "flashlib_cake_kmeans.launch_dispatch_kernel_0043", + "flashlib_cake_kmeans.launch_dispatch_kernel_0044", + "flashlib_cake_kmeans.launch_dispatch_kernel_0045", + "flashlib_cake_kmeans.launch_dispatch_kernel_0046", + "flashlib_cake_kmeans.launch_dispatch_kernel_0047", + "flashlib_cake_kmeans.launch_dispatch_kernel_0048", + "flashlib_cake_kmeans.launch_dispatch_kernel_0049", + "flashlib_cake_kmeans.launch_dispatch_kernel_0050", + "flashlib_cake_kmeans.launch_dispatch_kernel_0051", + "flashlib_cake_kmeans.launch_dispatch_kernel_0052", + "flashlib_cake_kmeans.launch_dispatch_kernel_0053", + "flashlib_cake_kmeans.launch_dispatch_kernel_0054", + "flashlib_cake_kmeans.launch_dispatch_kernel_0055", + "flashlib_cake_kmeans.launch_dispatch_kernel_0056", + "flashlib_cake_kmeans.launch_dispatch_kernel_0057", + "flashlib_cake_kmeans.launch_dispatch_kernel_0058", + "flashlib_cake_kmeans.launch_dispatch_kernel_0059", + "flashlib_cake_kmeans.launch_dispatch_kernel_0060", + "flashlib_cake_kmeans.launch_dispatch_kernel_0061", + "flashlib_cake_kmeans.launch_dispatch_kernel_0062", + "flashlib_cake_kmeans.launch_dispatch_kernel_0063", + "flashlib_cake_kmeans.launch_dispatch_kernel_0064", + "flashlib_cake_kmeans.launch_dispatch_kernel_0065", + "flashlib_cake_kmeans.launch_dispatch_kernel_0066", + "flashlib_cake_kmeans.launch_dispatch_kernel_0067", + "flashlib_cake_kmeans.launch_dispatch_kernel_0068", + "flashlib_cake_kmeans.launch_dispatch_kernel_0069", + "flashlib_cake_kmeans.launch_dispatch_kernel_0070", + "flashlib_cake_kmeans.launch_dispatch_kernel_0071", + "flashlib_cake_kmeans.launch_dispatch_kernel_0072", + "flashlib_cake_kmeans.launch_dispatch_kernel_0073", + "flashlib_cake_kmeans.launch_dispatch_kernel_0074", + "flashlib_cake_kmeans.launch_dispatch_kernel_0075", + "flashlib_cake_kmeans.launch_dispatch_kernel_0076", + "flashlib_cake_kmeans.launch_dispatch_kernel_0077", + "flashlib_cake_kmeans.launch_dispatch_kernel_0078", + "flashlib_cake_kmeans.launch_dispatch_kernel_0079", + "flashlib_cake_kmeans.launch_dispatch_kernel_0080", + "flashlib_cake_kmeans.launch_dispatch_kernel_0081", + "flashlib_cake_kmeans.launch_dispatch_kernel_0082", + "flashlib_cake_kmeans.launch_dispatch_kernel_0083", + "flashlib_cake_kmeans.launch_dispatch_kernel_0084", + "flashlib_cake_kmeans.launch_dispatch_kernel_0085", + "flashlib_cake_kmeans.launch_dispatch_kernel_0086", + "flashlib_cake_kmeans.launch_dispatch_kernel_0087", + "flashlib_cake_kmeans.launch_dispatch_kernel_0088", + "flashlib_cake_kmeans.launch_dispatch_kernel_0089", + "flashlib_cake_kmeans.launch_dispatch_kernel_0090", + "flashlib_cake_kmeans.launch_dispatch_kernel_0091", + "flashlib_cake_kmeans.launch_dispatch_kernel_0092", + "flashlib_cake_kmeans.launch_dispatch_kernel_0093", + "flashlib_cake_kmeans.launch_dispatch_kernel_0094", + "flashlib_cake_kmeans.launch_dispatch_kernel_0095", + "flashlib_cake_kmeans.launch_dispatch_kernel_0096", + "flashlib_cake_kmeans.launch_dispatch_kernel_0097", + "flashlib_cake_kmeans.launch_dispatch_kernel_0098", + "flashlib_cake_kmeans.launch_dispatch_kernel_0099", + "flashlib_cake_kmeans.launch_dispatch_kernel_0100", + "flashlib_cake_kmeans.launch_dispatch_kernel_0101", + "flashlib_cake_kmeans.launch_dispatch_kernel_0102", + "flashlib_cake_kmeans.launch_dispatch_kernel_0103", + "flashlib_cake_kmeans.launch_dispatch_kernel_0104", + "flashlib_cake_kmeans.launch_dispatch_kernel_0105", + "flashlib_cake_kmeans.launch_dispatch_kernel_0106", + "flashlib_cake_kmeans.launch_dispatch_kernel_0107", + "flashlib_cake_kmeans.launch_dispatch_kernel_0108", + "flashlib_cake_kmeans.launch_dispatch_kernel_0109", + "flashlib_cake_kmeans.launch_dispatch_kernel_0110", + "flashlib_cake_kmeans.launch_dispatch_kernel_0111", + "flashlib_cake_kmeans.launch_dispatch_kernel_0112", + "flashlib_cake_kmeans.launch_dispatch_kernel_0113", + "flashlib_cake_kmeans.launch_dispatch_kernel_0114", + "flashlib_cake_kmeans.launch_dispatch_kernel_0115", + "flashlib_cake_kmeans.launch_dispatch_kernel_0116", + "flashlib_cake_kmeans.launch_dispatch_kernel_0117", + "flashlib_cake_kmeans.launch_dispatch_kernel_0118", + "flashlib_cake_kmeans.launch_dispatch_kernel_0119", + "flashlib_cake_kmeans.launch_dispatch_kernel_0120", + "flashlib_cake_kmeans.launch_dispatch_kernel_0121", + "flashlib_cake_kmeans.launch_dispatch_kernel_0122", + "flashlib_cake_kmeans.launch_dispatch_kernel_0123", + "flashlib_cake_kmeans.launch_dispatch_kernel_0124", + "flashlib_cake_kmeans.launch_dispatch_kernel_0125", + "flashlib_cake_kmeans.launch_dispatch_kernel_0126", + "flashlib_cake_kmeans.launch_dispatch_kernel_0127", + "flashlib_cake_kmeans.launch_dispatch_kernel_0128", + "flashlib_cake_kmeans.launch_dispatch_kernel_0129", + "flashlib_cake_kmeans.launch_dispatch_kernel_0130", + "flashlib_cake_kmeans.launch_dispatch_kernel_0131", + "flashlib_cake_kmeans.launch_dispatch_kernel_0132", + "flashlib_cake_kmeans.launch_dispatch_kernel_0133", + "flashlib_cake_kmeans.launch_dispatch_kernel_0134", + "flashlib_cake_kmeans.launch_dispatch_kernel_0135", + "flashlib_cake_kmeans.launch_dispatch_kernel_0136", + "flashlib_cake_kmeans.launch_dispatch_kernel_0137", + "flashlib_cake_kmeans.launch_dispatch_kernel_0138", + "flashlib_cake_kmeans.launch_dispatch_kernel_0139", + "flashlib_cake_kmeans.launch_dispatch_kernel_0140", + "flashlib_cake_kmeans.launch_dispatch_kernel_0141", + "flashlib_cake_kmeans.launch_dispatch_kernel_0142", + "flashlib_cake_kmeans.launch_dispatch_kernel_0143", + "flashlib_cake_kmeans.launch_dispatch_kernel_0144", + "flashlib_cake_kmeans.launch_dispatch_kernel_0145", + "flashlib_cake_kmeans.launch_dispatch_kernel_0146", + "flashlib_cake_kmeans.launch_dispatch_kernel_0147", + "flashlib_cake_kmeans.launch_dispatch_kernel_0148", + "flashlib_cake_kmeans.launch_dispatch_kernel_0149", + "flashlib_cake_kmeans.launch_dispatch_kernel_0150", + "flashlib_cake_kmeans.launch_dispatch_kernel_0151", + "flashlib_cake_kmeans.launch_dispatch_kernel_0152", + "flashlib_cake_kmeans.launch_dispatch_kernel_0153", + "flashlib_cake_kmeans.launch_dispatch_kernel_0154", + "flashlib_cake_kmeans.launch_dispatch_kernel_0155", + "flashlib_cake_kmeans.launch_dispatch_kernel_0156", + "flashlib_cake_kmeans.launch_dispatch_kernel_0157", + "flashlib_cake_kmeans.launch_dispatch_kernel_0158", + "flashlib_cake_kmeans.launch_dispatch_kernel_0159", + "flashlib_cake_kmeans.launch_dispatch_kernel_0160", + "flashlib_cake_kmeans.launch_dispatch_kernel_0161", + "flashlib_cake_kmeans.launch_dispatch_kernel_0162", + "flashlib_cake_kmeans.launch_dispatch_kernel_0163", + "flashlib_cake_kmeans.launch_dispatch_kernel_0164", + "flashlib_cake_kmeans.launch_dispatch_kernel_0165", + "flashlib_cake_kmeans.launch_dispatch_kernel_0166", + "flashlib_cake_kmeans.launch_dispatch_kernel_0167", + "flashlib_cake_kmeans.launch_dispatch_kernel_0168", + "flashlib_cake_kmeans.launch_dispatch_kernel_0169", + "flashlib_cake_kmeans.launch_dispatch_kernel_0170", + "flashlib_cake_kmeans.launch_dispatch_kernel_0171", + "flashlib_cake_kmeans.launch_dispatch_kernel_0172", + "flashlib_cake_kmeans.launch_dispatch_kernel_0173", + "flashlib_cake_kmeans.launch_dispatch_kernel_0174", + "flashlib_cake_kmeans.launch_dispatch_kernel_0175", + "flashlib_cake_kmeans.launch_dispatch_kernel_0176", + "flashlib_cake_kmeans.launch_dispatch_kernel_0177", + "flashlib_cake_kmeans.launch_dispatch_kernel_0178", + "flashlib_cake_kmeans.launch_dispatch_kernel_0179", + "flashlib_cake_kmeans.launch_dispatch_kernel_0180", + "flashlib_cake_kmeans.launch_dispatch_kernel_0181", + "flashlib_cake_kmeans.launch_dispatch_kernel_0182", + "flashlib_cake_kmeans.launch_dispatch_kernel_0183", + "flashlib_cake_kmeans.launch_dispatch_kernel_0184", + "flashlib_cake_kmeans.launch_dispatch_kernel_0185", + "flashlib_cake_kmeans.launch_dispatch_kernel_0186", + "flashlib_cake_kmeans.launch_dispatch_kernel_0187", + "flashlib_cake_kmeans.launch_dispatch_kernel_0188", + "flashlib_cake_kmeans.launch_dispatch_kernel_0189", + "flashlib_cake_kmeans.launch_dispatch_kernel_0190", + "flashlib_cake_kmeans.launch_dispatch_kernel_0191", + "flashlib_cake_kmeans.launch_dispatch_kernel_0192", + "flashlib_cake_kmeans.launch_dispatch_kernel_0193", + "flashlib_cake_kmeans.launch_dispatch_kernel_0194", + "flashlib_cake_kmeans.launch_dispatch_kernel_0195", + "flashlib_cake_kmeans.launch_dispatch_kernel_0196", + "flashlib_cake_kmeans.launch_dispatch_kernel_0197", + "flashlib_cake_kmeans.launch_dispatch_kernel_0198", + "flashlib_cake_kmeans.launch_dispatch_kernel_0199", + "flashlib_cake_kmeans.launch_dispatch_kernel_0200", + "flashlib_cake_kmeans.launch_dispatch_kernel_0201", + "flashlib_cake_kmeans.launch_dispatch_kernel_0202", + "flashlib_cake_kmeans.launch_dispatch_kernel_0203", + "flashlib_cake_kmeans.launch_dispatch_kernel_0204", + "flashlib_cake_kmeans.launch_dispatch_kernel_0205", + "flashlib_cake_kmeans.launch_dispatch_kernel_0206", + "flashlib_cake_kmeans.launch_dispatch_kernel_0207", + "flashlib_cake_kmeans.launch_dispatch_kernel_0208", + "flashlib_cake_kmeans.launch_dispatch_kernel_0209", + "flashlib_cake_kmeans.launch_dispatch_kernel_0210", + "flashlib_cake_kmeans.launch_dispatch_kernel_0211", + "flashlib_cake_kmeans.launch_dispatch_kernel_0212", + "flashlib_cake_kmeans.launch_dispatch_kernel_0213", + "flashlib_cake_kmeans.launch_dispatch_kernel_0214", + "flashlib_cake_kmeans.launch_dispatch_kernel_0215", + "flashlib_cake_kmeans.launch_dispatch_kernel_0216", + "flashlib_cake_kmeans.launch_dispatch_kernel_0217", + "flashlib_cake_kmeans.launch_dispatch_kernel_0218", + "flashlib_cake_kmeans.launch_dispatch_kernel_0219", + "flashlib_cake_kmeans.launch_dispatch_kernel_0220", + "flashlib_cake_kmeans.launch_dispatch_kernel_0221", + "flashlib_cake_kmeans.launch_dispatch_kernel_0222", + "flashlib_cake_kmeans.launch_dispatch_kernel_0223", + "flashlib_cake_kmeans.launch_dispatch_kernel_0224", + "flashlib_cake_kmeans.launch_dispatch_kernel_0225", + "flashlib_cake_kmeans.launch_dispatch_kernel_0226", + "flashlib_cake_kmeans.launch_dispatch_kernel_0227", + "flashlib_cake_kmeans.launch_dispatch_kernel_0228", + "flashlib_cake_kmeans.launch_dispatch_kernel_0229", + "flashlib_cake_kmeans.launch_dispatch_kernel_0230", + "flashlib_cake_kmeans.launch_dispatch_kernel_0231", + "flashlib_cake_kmeans.launch_dispatch_kernel_0232", + "flashlib_cake_kmeans.launch_dispatch_kernel_0233", + "flashlib_cake_kmeans.launch_dispatch_kernel_0234", + "flashlib_cake_kmeans.launch_dispatch_kernel_0235", + "flashlib_cake_kmeans.launch_dispatch_kernel_0236", + "flashlib_cake_kmeans.launch_dispatch_kernel_0237", + "flashlib_cake_kmeans.launch_dispatch_kernel_0238", + "flashlib_cake_kmeans.launch_dispatch_kernel_0239", + "flashlib_cake_kmeans.launch_dispatch_kernel_0240", + "flashlib_cake_kmeans.launch_dispatch_kernel_0241", + "flashlib_cake_kmeans.launch_dispatch_kernel_0242", + "flashlib_cake_kmeans.launch_dispatch_kernel_0243", + "flashlib_cake_kmeans.launch_dispatch_kernel_0244", + "flashlib_cake_kmeans.launch_dispatch_kernel_0245", + "flashlib_cake_kmeans.launch_dispatch_kernel_0246", + "flashlib_cake_kmeans.launch_dispatch_kernel_0247", + "flashlib_cake_kmeans.launch_dispatch_kernel_0248", + "flashlib_cake_kmeans.launch_dispatch_kernel_0249", + "flashlib_cake_kmeans.launch_dispatch_kernel_0250", + "flashlib_cake_kmeans.launch_dispatch_kernel_0251", + "flashlib_cake_kmeans.launch_dispatch_kernel_0252", + "flashlib_cake_kmeans.launch_dispatch_kernel_0253", + "flashlib_cake_kmeans.launch_dispatch_kernel_0254", + "flashlib_cake_kmeans.launch_dispatch_kernel_0255", + "flashlib_cake_kmeans.launch_dispatch_kernel_0256", + "flashlib_cake_kmeans.launch_dispatch_kernel_0257", + "flashlib_cake_kmeans.launch_dispatch_kernel_0258", + "flashlib_cake_kmeans.launch_dispatch_kernel_0259", + "flashlib_cake_kmeans.launch_dispatch_kernel_0260", + "flashlib_cake_kmeans.launch_dispatch_kernel_0261", + "flashlib_cake_kmeans.launch_dispatch_kernel_0262", + "flashlib_cake_kmeans.launch_dispatch_kernel_0263", + "flashlib_cake_kmeans.launch_dispatch_kernel_0264", + "flashlib_cake_kmeans.launch_dispatch_kernel_0265", + "flashlib_cake_kmeans.launch_dispatch_kernel_0266", + "flashlib_cake_kmeans.launch_dispatch_kernel_0267", + "flashlib_cake_kmeans.launch_dispatch_kernel_0268", + "flashlib_cake_kmeans.launch_dispatch_kernel_0269", + "flashlib_cake_kmeans.launch_dispatch_kernel_0270", + "flashlib_cake_kmeans.launch_dispatch_kernel_0271", + "flashlib_cake_kmeans.launch_dispatch_kernel_0272", + "flashlib_cake_kmeans.launch_dispatch_kernel_0273", + "flashlib_cake_kmeans.launch_dispatch_kernel_0274", + "flashlib_cake_kmeans.launch_dispatch_kernel_0275", + "flashlib_cake_kmeans.launch_dispatch_kernel_0276", + "flashlib_cake_kmeans.launch_dispatch_kernel_0277", + "flashlib_cake_kmeans.launch_dispatch_kernel_0278", + "flashlib_cake_kmeans.launch_dispatch_kernel_0279", + "flashlib_cake_kmeans.launch_dispatch_kernel_0280", + "flashlib_cake_kmeans.launch_dispatch_kernel_0281", + "flashlib_cake_kmeans.launch_dispatch_kernel_0282", + "flashlib_cake_kmeans.launch_dispatch_kernel_0283", + "flashlib_cake_kmeans.launch_dispatch_kernel_0284", + "flashlib_cake_kmeans.launch_dispatch_kernel_0285", + "flashlib_cake_kmeans.launch_dispatch_kernel_0286", + "flashlib_cake_kmeans.launch_dispatch_kernel_0287", + "flashlib_cake_kmeans.launch_dispatch_kernel_0288", + "flashlib_cake_kmeans.launch_dispatch_kernel_0289", + "flashlib_cake_kmeans.launch_dispatch_kernel_0290", + "flashlib_cake_kmeans.launch_dispatch_kernel_0291", + "flashlib_cake_kmeans.launch_dispatch_kernel_0292", + "flashlib_cake_kmeans.launch_dispatch_kernel_0293", + "flashlib_cake_kmeans.launch_dispatch_kernel_0294", + "flashlib_cake_kmeans.launch_dispatch_kernel_0295", + "flashlib_cake_kmeans.launch_dispatch_kernel_0296", + "flashlib_cake_kmeans.launch_dispatch_kernel_0297", + "flashlib_cake_kmeans.launch_dispatch_kernel_0298", + "flashlib_cake_kmeans.launch_dispatch_kernel_0299", + "flashlib_cake_kmeans.launch_dispatch_kernel_0300", + "flashlib_cake_kmeans.launch_dispatch_kernel_0301", + "flashlib_cake_kmeans.launch_dispatch_kernel_0302", + "flashlib_cake_kmeans.launch_dispatch_kernel_0303", + "flashlib_cake_kmeans.launch_dispatch_kernel_0304", + "flashlib_cake_kmeans.launch_dispatch_kernel_0305", + "flashlib_cake_kmeans.launch_dispatch_kernel_0306", + "flashlib_cake_kmeans.launch_dispatch_kernel_0307", + "flashlib_cake_kmeans.launch_dispatch_kernel_0308", + "flashlib_cake_kmeans.launch_dispatch_kernel_0309", + "flashlib_cake_kmeans.launch_dispatch_kernel_0310", + "flashlib_cake_kmeans.launch_dispatch_kernel_0311", + "flashlib_cake_kmeans.launch_dispatch_kernel_0312", + "flashlib_cake_kmeans.launch_dispatch_kernel_0313", + "flashlib_cake_kmeans.launch_dispatch_kernel_0314", + "flashlib_cake_kmeans.launch_dispatch_kernel_0315", + "flashlib_cake_kmeans.launch_dispatch_kernel_0316", + "flashlib_cake_kmeans.launch_dispatch_kernel_0317", + "flashlib_cake_kmeans.launch_dispatch_kernel_0318", + "flashlib_cake_kmeans.launch_dispatch_kernel_0319", + "flashlib_cake_kmeans.launch_dispatch_kernel_0320", + "flashlib_cake_kmeans.launch_dispatch_kernel_0321", + "flashlib_cake_kmeans.launch_dispatch_kernel_0322", + "flashlib_cake_kmeans.launch_dispatch_kernel_0323", + "flashlib_cake_kmeans.launch_dispatch_kernel_0324", + "flashlib_cake_kmeans.launch_dispatch_kernel_0325", + "flashlib_cake_kmeans.launch_dispatch_kernel_0326", + "flashlib_cake_kmeans.launch_dispatch_kernel_0327", + "flashlib_cake_kmeans.launch_dispatch_kernel_0328", + "flashlib_cake_kmeans.launch_dispatch_kernel_0329", + "flashlib_cake_kmeans.launch_dispatch_kernel_0330", + "flashlib_cake_kmeans.launch_dispatch_kernel_0331", + "flashlib_cake_kmeans.launch_dispatch_kernel_0332", + "flashlib_cake_kmeans.launch_dispatch_kernel_0333", + "flashlib_cake_kmeans.launch_dispatch_kernel_0334", + "flashlib_cake_kmeans.launch_dispatch_kernel_0335", + "flashlib_cake_kmeans.launch_dispatch_kernel_0336", + "flashlib_cake_kmeans.launch_dispatch_kernel_0337", + "flashlib_cake_kmeans.launch_dispatch_kernel_0338", + "flashlib_cake_kmeans.launch_dispatch_kernel_0339", + "flashlib_cake_kmeans.launch_dispatch_kernel_0340", + "flashlib_cake_kmeans.launch_dispatch_kernel_0341", + "flashlib_cake_kmeans.launch_dispatch_kernel_0342", + "flashlib_cake_kmeans.launch_dispatch_kernel_0343", + "flashlib_cake_kmeans.launch_dispatch_kernel_0344", + "flashlib_cake_kmeans.launch_dispatch_kernel_0345", + "flashlib_cake_kmeans.launch_dispatch_kernel_0346", + "flashlib_cake_kmeans.launch_dispatch_kernel_0347", + "flashlib_cake_kmeans.launch_dispatch_kernel_0348", + "flashlib_cake_kmeans.launch_dispatch_kernel_0349", + "flashlib_cake_kmeans.launch_dispatch_kernel_0350", + "flashlib_cake_kmeans.launch_dispatch_kernel_0351", + "flashlib_cake_kmeans.launch_dispatch_kernel_0352", + "flashlib_cake_kmeans.launch_dispatch_kernel_0353", + "flashlib_cake_kmeans.launch_dispatch_kernel_0354", + "flashlib_cake_kmeans.launch_dispatch_kernel_0355", + "flashlib_cake_kmeans.launch_dispatch_kernel_0356", + "flashlib_cake_kmeans.launch_dispatch_kernel_0357" + ], + "tensor_interop": "dlpack-zero-copy" + } +} diff --git a/cake_exports/kmeans/src/flashlib_cake_kmeans/tvm_ffi.py b/cake_exports/kmeans/src/flashlib_cake_kmeans/tvm_ffi.py new file mode 100644 index 00000000..8a639de1 --- /dev/null +++ b/cake_exports/kmeans/src/flashlib_cake_kmeans/tvm_ffi.py @@ -0,0 +1,179 @@ +"""Optional Apache TVM FFI adapters for the exported kernel package.""" + +from __future__ import annotations + +import contextlib +import importlib +import json +from pathlib import Path +from typing import Any + +from .kernels import get_kernel + + +_PACKAGE = 'flashlib_cake_kmeans' +_KERNEL_ALIASES = ['dispatch_kernel_0000', 'dispatch_kernel_0001', 'dispatch_kernel_0002', 'dispatch_kernel_0003', 'dispatch_kernel_0004', 'dispatch_kernel_0005', 'dispatch_kernel_0006', 'dispatch_kernel_0007', 'dispatch_kernel_0008', 'dispatch_kernel_0009', 'dispatch_kernel_0010', 'dispatch_kernel_0011', 'dispatch_kernel_0012', 'dispatch_kernel_0013', 'dispatch_kernel_0014', 'dispatch_kernel_0015', 'dispatch_kernel_0016', 'dispatch_kernel_0017', 'dispatch_kernel_0018', 'dispatch_kernel_0019', 'dispatch_kernel_0020', 'dispatch_kernel_0021', 'dispatch_kernel_0022', 'dispatch_kernel_0023', 'dispatch_kernel_0024', 'dispatch_kernel_0025', 'dispatch_kernel_0026', 'dispatch_kernel_0027', 'dispatch_kernel_0028', 'dispatch_kernel_0029', 'dispatch_kernel_0030', 'dispatch_kernel_0031', 'dispatch_kernel_0032', 'dispatch_kernel_0033', 'dispatch_kernel_0034', 'dispatch_kernel_0035', 'dispatch_kernel_0036', 'dispatch_kernel_0037', 'dispatch_kernel_0038', 'dispatch_kernel_0039', 'dispatch_kernel_0040', 'dispatch_kernel_0041', 'dispatch_kernel_0042', 'dispatch_kernel_0043', 'dispatch_kernel_0044', 'dispatch_kernel_0045', 'dispatch_kernel_0046', 'dispatch_kernel_0047', 'dispatch_kernel_0048', 'dispatch_kernel_0049', 'dispatch_kernel_0050', 'dispatch_kernel_0051', 'dispatch_kernel_0052', 'dispatch_kernel_0053', 'dispatch_kernel_0054', 'dispatch_kernel_0055', 'dispatch_kernel_0056', 'dispatch_kernel_0057', 'dispatch_kernel_0058', 'dispatch_kernel_0059', 'dispatch_kernel_0060', 'dispatch_kernel_0061', 'dispatch_kernel_0062', 'dispatch_kernel_0063', 'dispatch_kernel_0064', 'dispatch_kernel_0065', 'dispatch_kernel_0066', 'dispatch_kernel_0067', 'dispatch_kernel_0068', 'dispatch_kernel_0069', 'dispatch_kernel_0070', 'dispatch_kernel_0071', 'dispatch_kernel_0072', 'dispatch_kernel_0073', 'dispatch_kernel_0074', 'dispatch_kernel_0075', 'dispatch_kernel_0076', 'dispatch_kernel_0077', 'dispatch_kernel_0078', 'dispatch_kernel_0079', 'dispatch_kernel_0080', 'dispatch_kernel_0081', 'dispatch_kernel_0082', 'dispatch_kernel_0083', 'dispatch_kernel_0084', 'dispatch_kernel_0085', 'dispatch_kernel_0086', 'dispatch_kernel_0087', 'dispatch_kernel_0088', 'dispatch_kernel_0089', 'dispatch_kernel_0090', 'dispatch_kernel_0091', 'dispatch_kernel_0092', 'dispatch_kernel_0093', 'dispatch_kernel_0094', 'dispatch_kernel_0095', 'dispatch_kernel_0096', 'dispatch_kernel_0097', 'dispatch_kernel_0098', 'dispatch_kernel_0099', 'dispatch_kernel_0100', 'dispatch_kernel_0101', 'dispatch_kernel_0102', 'dispatch_kernel_0103', 'dispatch_kernel_0104', 'dispatch_kernel_0105', 'dispatch_kernel_0106', 'dispatch_kernel_0107', 'dispatch_kernel_0108', 'dispatch_kernel_0109', 'dispatch_kernel_0110', 'dispatch_kernel_0111', 'dispatch_kernel_0112', 'dispatch_kernel_0113', 'dispatch_kernel_0114', 'dispatch_kernel_0115', 'dispatch_kernel_0116', 'dispatch_kernel_0117', 'dispatch_kernel_0118', 'dispatch_kernel_0119', 'dispatch_kernel_0120', 'dispatch_kernel_0121', 'dispatch_kernel_0122', 'dispatch_kernel_0123', 'dispatch_kernel_0124', 'dispatch_kernel_0125', 'dispatch_kernel_0126', 'dispatch_kernel_0127', 'dispatch_kernel_0128', 'dispatch_kernel_0129', 'dispatch_kernel_0130', 'dispatch_kernel_0131', 'dispatch_kernel_0132', 'dispatch_kernel_0133', 'dispatch_kernel_0134', 'dispatch_kernel_0135', 'dispatch_kernel_0136', 'dispatch_kernel_0137', 'dispatch_kernel_0138', 'dispatch_kernel_0139', 'dispatch_kernel_0140', 'dispatch_kernel_0141', 'dispatch_kernel_0142', 'dispatch_kernel_0143', 'dispatch_kernel_0144', 'dispatch_kernel_0145', 'dispatch_kernel_0146', 'dispatch_kernel_0147', 'dispatch_kernel_0148', 'dispatch_kernel_0149', 'dispatch_kernel_0150', 'dispatch_kernel_0151', 'dispatch_kernel_0152', 'dispatch_kernel_0153', 'dispatch_kernel_0154', 'dispatch_kernel_0155', 'dispatch_kernel_0156', 'dispatch_kernel_0157', 'dispatch_kernel_0158', 'dispatch_kernel_0159', 'dispatch_kernel_0160', 'dispatch_kernel_0161', 'dispatch_kernel_0162', 'dispatch_kernel_0163', 'dispatch_kernel_0164', 'dispatch_kernel_0165', 'dispatch_kernel_0166', 'dispatch_kernel_0167', 'dispatch_kernel_0168', 'dispatch_kernel_0169', 'dispatch_kernel_0170', 'dispatch_kernel_0171', 'dispatch_kernel_0172', 'dispatch_kernel_0173', 'dispatch_kernel_0174', 'dispatch_kernel_0175', 'dispatch_kernel_0176', 'dispatch_kernel_0177', 'dispatch_kernel_0178', 'dispatch_kernel_0179', 'dispatch_kernel_0180', 'dispatch_kernel_0181', 'dispatch_kernel_0182', 'dispatch_kernel_0183', 'dispatch_kernel_0184', 'dispatch_kernel_0185', 'dispatch_kernel_0186', 'dispatch_kernel_0187', 'dispatch_kernel_0188', 'dispatch_kernel_0189', 'dispatch_kernel_0190', 'dispatch_kernel_0191', 'dispatch_kernel_0192', 'dispatch_kernel_0193', 'dispatch_kernel_0194', 'dispatch_kernel_0195', 'dispatch_kernel_0196', 'dispatch_kernel_0197', 'dispatch_kernel_0198', 'dispatch_kernel_0199', 'dispatch_kernel_0200', 'dispatch_kernel_0201', 'dispatch_kernel_0202', 'dispatch_kernel_0203', 'dispatch_kernel_0204', 'dispatch_kernel_0205', 'dispatch_kernel_0206', 'dispatch_kernel_0207', 'dispatch_kernel_0208', 'dispatch_kernel_0209', 'dispatch_kernel_0210', 'dispatch_kernel_0211', 'dispatch_kernel_0212', 'dispatch_kernel_0213', 'dispatch_kernel_0214', 'dispatch_kernel_0215', 'dispatch_kernel_0216', 'dispatch_kernel_0217', 'dispatch_kernel_0218', 'dispatch_kernel_0219', 'dispatch_kernel_0220', 'dispatch_kernel_0221', 'dispatch_kernel_0222', 'dispatch_kernel_0223', 'dispatch_kernel_0224', 'dispatch_kernel_0225', 'dispatch_kernel_0226', 'dispatch_kernel_0227', 'dispatch_kernel_0228', 'dispatch_kernel_0229', 'dispatch_kernel_0230', 'dispatch_kernel_0231', 'dispatch_kernel_0232', 'dispatch_kernel_0233', 'dispatch_kernel_0234', 'dispatch_kernel_0235', 'dispatch_kernel_0236', 'dispatch_kernel_0237', 'dispatch_kernel_0238', 'dispatch_kernel_0239', 'dispatch_kernel_0240', 'dispatch_kernel_0241', 'dispatch_kernel_0242', 'dispatch_kernel_0243', 'dispatch_kernel_0244', 'dispatch_kernel_0245', 'dispatch_kernel_0246', 'dispatch_kernel_0247', 'dispatch_kernel_0248', 'dispatch_kernel_0249', 'dispatch_kernel_0250', 'dispatch_kernel_0251', 'dispatch_kernel_0252', 'dispatch_kernel_0253', 'dispatch_kernel_0254', 'dispatch_kernel_0255', 'dispatch_kernel_0256', 'dispatch_kernel_0257', 'dispatch_kernel_0258', 'dispatch_kernel_0259', 'dispatch_kernel_0260', 'dispatch_kernel_0261', 'dispatch_kernel_0262', 'dispatch_kernel_0263', 'dispatch_kernel_0264', 'dispatch_kernel_0265', 'dispatch_kernel_0266', 'dispatch_kernel_0267', 'dispatch_kernel_0268', 'dispatch_kernel_0269', 'dispatch_kernel_0270', 'dispatch_kernel_0271', 'dispatch_kernel_0272', 'dispatch_kernel_0273', 'dispatch_kernel_0274', 'dispatch_kernel_0275', 'dispatch_kernel_0276', 'dispatch_kernel_0277', 'dispatch_kernel_0278', 'dispatch_kernel_0279', 'dispatch_kernel_0280', 'dispatch_kernel_0281', 'dispatch_kernel_0282', 'dispatch_kernel_0283', 'dispatch_kernel_0284', 'dispatch_kernel_0285', 'dispatch_kernel_0286', 'dispatch_kernel_0287', 'dispatch_kernel_0288', 'dispatch_kernel_0289', 'dispatch_kernel_0290', 'dispatch_kernel_0291', 'dispatch_kernel_0292', 'dispatch_kernel_0293', 'dispatch_kernel_0294', 'dispatch_kernel_0295', 'dispatch_kernel_0296', 'dispatch_kernel_0297', 'dispatch_kernel_0298', 'dispatch_kernel_0299', 'dispatch_kernel_0300', 'dispatch_kernel_0301', 'dispatch_kernel_0302', 'dispatch_kernel_0303', 'dispatch_kernel_0304', 'dispatch_kernel_0305', 'dispatch_kernel_0306', 'dispatch_kernel_0307', 'dispatch_kernel_0308', 'dispatch_kernel_0309', 'dispatch_kernel_0310', 'dispatch_kernel_0311', 'dispatch_kernel_0312', 'dispatch_kernel_0313', 'dispatch_kernel_0314', 'dispatch_kernel_0315', 'dispatch_kernel_0316', 'dispatch_kernel_0317', 'dispatch_kernel_0318', 'dispatch_kernel_0319', 'dispatch_kernel_0320', 'dispatch_kernel_0321', 'dispatch_kernel_0322', 'dispatch_kernel_0323', 'dispatch_kernel_0324', 'dispatch_kernel_0325', 'dispatch_kernel_0326', 'dispatch_kernel_0327', 'dispatch_kernel_0328', 'dispatch_kernel_0329', 'dispatch_kernel_0330', 'dispatch_kernel_0331', 'dispatch_kernel_0332', 'dispatch_kernel_0333', 'dispatch_kernel_0334', 'dispatch_kernel_0335', 'dispatch_kernel_0336', 'dispatch_kernel_0337', 'dispatch_kernel_0338', 'dispatch_kernel_0339', 'dispatch_kernel_0340', 'dispatch_kernel_0341', 'dispatch_kernel_0342', 'dispatch_kernel_0343', 'dispatch_kernel_0344', 'dispatch_kernel_0345', 'dispatch_kernel_0346', 'dispatch_kernel_0347', 'dispatch_kernel_0348', 'dispatch_kernel_0349', 'dispatch_kernel_0350', 'dispatch_kernel_0351', 'dispatch_kernel_0352', 'dispatch_kernel_0353', 'dispatch_kernel_0354', 'dispatch_kernel_0355', 'dispatch_kernel_0356', 'dispatch_kernel_0357'] +_REGISTERED: dict[str, tuple[str, ...]] = {} + + +class _RawCUDAStream: + def __init__(self, handle: int): + self.cuda_stream = int(handle) + + +def _manifest() -> dict[str, Any]: + return json.loads(Path(__file__).with_name("manifest.json").read_text(encoding="utf-8")) + + +def _planned_public_exports() -> dict[str, str] | None: + export_plan = _manifest().get("export_plan", {}) + if "tvm_ffi_exports" in export_plan: + return dict(export_plan["tvm_ffi_exports"]) + if "package_exports" in export_plan: + return dict(export_plan["package_exports"]) + return None + + +def _public_export_names() -> tuple[str, ...]: + planned = _planned_public_exports() + if planned is not None: + return tuple(planned) + package = importlib.import_module(__package__) + excluded = { + "KERNELS", + "ExportedKernel", + "get_kernel", + "register_tvm_ffi", + "tvm_ffi_function_names", + *(_KERNEL_ALIASES), + *(f"launch_{name}" for name in _KERNEL_ALIASES), + } + return tuple( + name + for name in getattr(package, "__all__", ()) + if name not in excluded and callable(getattr(package, name, None)) + ) + + +def tvm_ffi_function_names(namespace: str | None = None) -> tuple[str, ...]: + """Return the deterministic TVM FFI global names this package registers.""" + + prefix = namespace or _PACKAGE + public_names = [f"{prefix}.{name}" for name in _public_export_names()] + kernel_names = [f"{prefix}.launch_{name}" for name in _KERNEL_ALIASES] + return tuple(public_names + kernel_names) + + +def _tensor_stream(arg: Any, tvm_ffi: Any) -> tuple[int, int] | None: + if not isinstance(arg, tvm_ffi.Tensor): + return None + device = arg.device + device_type = getattr(device, "type", None) or str(device).split(":", 1)[0] + if str(device_type) != "cuda": + return None + device_id = int(getattr(device, "index", 0)) + return device_id, int(tvm_ffi.get_raw_stream(device)) + + +def _convert_arg(arg: Any, tvm_ffi: Any) -> Any: + if not isinstance(arg, tvm_ffi.Tensor): + return arg + import torch + + return torch.from_dlpack(arg) + + +def _torch_stream_context(stream: tuple[int, int] | None): + if stream is None: + return contextlib.nullcontext() + import torch + + device_id, handle = stream + return torch.cuda.stream(torch.cuda.ExternalStream(handle, device=device_id)) + + +def _semantic_target(public_name: str): + planned = _planned_public_exports() + target = None if planned is None else planned.get(public_name) + if target is None: + return getattr(importlib.import_module(__package__), public_name) + module_name, separator, attr = target.partition(":") + if not separator: + raise ValueError(f"invalid package export target: {target!r}") + module = importlib.import_module(module_name, package=__package__) + return getattr(module, attr) + + +def _semantic_wrapper(public_name: str, tvm_ffi: Any): + function = _semantic_target(public_name) + + def call(*args): + stream = next((item for arg in args if (item := _tensor_stream(arg, tvm_ffi)) is not None), None) + converted = tuple(_convert_arg(arg, tvm_ffi) for arg in args) + with _torch_stream_context(stream): + return function(*converted) + + return call + + +def _kernel_wrapper(alias: str, tvm_ffi: Any): + kernel = get_kernel(alias) + parameter_count = len(kernel.parameters) + + def call(*args): + expected = parameter_count + 7 + if len(args) != expected: + raise TypeError( + f"{alias} TVM FFI launch expects {expected} positional arguments " + f"({parameter_count} kernel + grid xyz + block xyz + shared_mem), got {len(args)}" + ) + kernel_args = args[:parameter_count] + config = tuple(int(value) for value in args[parameter_count:]) + grid = config[:3] + block = config[3:6] + shared_mem = config[6] + stream = next( + (item for arg in kernel_args if (item := _tensor_stream(arg, tvm_ffi)) is not None), + None, + ) + converted = tuple(_convert_arg(arg, tvm_ffi) for arg in kernel_args) + return kernel.launch( + *converted, + grid=grid, + block=block, + shared_mem=shared_mem, + stream=_RawCUDAStream(stream[1]) if stream is not None else None, + ) + + return call + + +def register_tvm_ffi(namespace: str | None = None, *, override: bool = False) -> tuple[str, ...]: + """Register semantic and low-level launch functions in Apache TVM FFI.""" + + try: + import tvm_ffi + except ImportError as exc: + raise ImportError( + 'TVM FFI support requires `python -m pip install -e ".[tvm-ffi]"`' + ) from exc + + prefix = namespace or _PACKAGE + if prefix in _REGISTERED and not override: + return _REGISTERED[prefix] + + registered: list[str] = [] + for public_name in _public_export_names(): + name = f"{prefix}.{public_name}" + tvm_ffi.register_global_func( + name, _semantic_wrapper(public_name, tvm_ffi), override=override + ) + registered.append(name) + for alias in _KERNEL_ALIASES: + name = f"{prefix}.launch_{alias}" + tvm_ffi.register_global_func(name, _kernel_wrapper(alias, tvm_ffi), override=override) + registered.append(name) + result = tuple(registered) + _REGISTERED[prefix] = result + return result + diff --git a/cake_exports/kmeans/tests/test_benchmark_harness.py b/cake_exports/kmeans/tests/test_benchmark_harness.py new file mode 100644 index 00000000..76ca7724 --- /dev/null +++ b/cake_exports/kmeans/tests/test_benchmark_harness.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[1] + + +def test_benchmark_harness_metadata_mode(tmp_path): + output = tmp_path / "benchmark_metadata.json" + env = os.environ.copy() + env["PYTHONPATH"] = str(ROOT / "src") + result = subprocess.run( + [ + sys.executable, + "benchmarks/benchmark_exported_kernels.py", + "--metadata-only", + "--json", + str(output), + ], + cwd=ROOT, + env=env, + check=False, + text=True, + capture_output=True, + ) + + assert result.returncode == 0, result.stdout + result.stderr + payload = json.loads(output.read_text(encoding="utf-8")) + assert payload["benchmark"] == "exported_kernel_compile" + assert payload["metadata_only"] is True + assert payload["kernels"] + assert payload["summary"]["kernel_count"] == len(payload["kernels"]) + + +def test_shape_benchmark_metadata_mode(tmp_path): + output = tmp_path / "shape_benchmark_metadata.json" + env = os.environ.copy() + env["PYTHONPATH"] = str(ROOT / "src") + result = subprocess.run( + [ + sys.executable, + "benchmarks/benchmark_shapes.py", + "--metadata-only", + "--json", + str(output), + ], + cwd=ROOT, + env=env, + check=False, + text=True, + capture_output=True, + ) + + assert result.returncode == 0, result.stdout + result.stderr + payload = json.loads(output.read_text(encoding="utf-8")) + assert payload["benchmark"] == "exported_kernel_shapes" + assert payload["metadata_only"] is True + assert payload["timing_backend_requested"] == "cupti" + diff --git a/cake_exports/kmeans/tests/test_correctness.py b/cake_exports/kmeans/tests/test_correctness.py new file mode 100644 index 00000000..23b5bfed --- /dev/null +++ b/cake_exports/kmeans/tests/test_correctness.py @@ -0,0 +1,255 @@ +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parents[1] +BENCHMARKS = ROOT / "benchmarks" +if str(BENCHMARKS) not in sys.path: + sys.path.insert(0, str(BENCHMARKS)) + + +def _benchmark_module(): + spec = importlib.util.spec_from_file_location("kmeans_benchmark", BENCHMARKS / "benchmark.py") + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +BENCHMARK = _benchmark_module() + + +def test_kmeans_runtime_api_is_exported(): + from flashlib_cake_kmeans import FlashKMeansAssignRuntime, init + + assert FlashKMeansAssignRuntime is not None + assert callable(init) + + +@pytest.mark.export_validation_shape +@pytest.mark.parametrize("row_index", range(len(BENCHMARK.FLASH_KMEANS_REGISTRY_SHAPES))) +def test_kmeans_matches_reference(row_index: int): + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("CUDA GPU required for exported-kernel correctness") + row = BENCHMARK.FLASH_KMEANS_REGISTRY_SHAPES[row_index] + from flashlib_cake_kmeans import init + + runtime = init() + baseline_adapter = BENCHMARK.TritonH20007cfRawAdapter( + device_index=runtime.device_index, + arch=runtime.arch, + ) + try: + result = BENCHMARK._run_shape( + row, + runtime=runtime, + baseline_adapter=baseline_adapter, + arch=None, + correctness=True, + benchmark=False, + reference_chunk_rows=128, + measurement_session_id="standalone-correctness-test", + measurement_order_seed=BENCHMARK.DEFAULT_MEASUREMENT_ORDER_SEED, + ) + finally: + runtime.clear() + baseline_adapter.clear() + assert result["route_matches_expected"], result + assert result["correct"], result + + +@pytest.mark.gpu +def test_kmeans_runtime_reuses_shape_after_intervening_shape(): + """Exercise A -> B -> fresh-pointer A on one live multi-shape runtime.""" + + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("CUDA GPU required for exported-kernel correctness") + from flashlib_cake_kmeans import init + + rows = {row["label"]: row for row in BENCHMARK.FLASH_KMEANS_SHAPES} + shape_a = rows["boundary_b1_n128_k256_d128"] + shape_b = rows["small_b1_n256_k256_d128"] + runtime = init() + + def run(row, *, variant: int): + x, centroids = BENCHMARK._make_inputs(row, variant=variant) + out = torch.empty((int(row["B"]), int(row["N"])), dtype=torch.int32, device=x.device) + result, info = runtime.compute( + x, + centroids, + out=out, + return_info=True, + ) + reference = BENCHMARK._reference_assign(x, centroids, chunk_rows=128) + torch.cuda.synchronize() + correctness = BENCHMARK._assignment_correctness(result, reference, x, centroids) + assert correctness["correct"], correctness + assert info["norm_mode"] == "fused_bf16_pair_row_norm:x_sq,c_sq" + assert info["norm_compute_fields"] == ["x_sq", "c_sq"] + return x, info + + try: + first_x, first_info = run(shape_a, variant=101) + _, middle_info = run(shape_b, variant=201) + fresh_x, revisit_info = run(shape_a, variant=102) + assert first_x.data_ptr() != fresh_x.data_ptr() + assert first_info["runtime_cache_hit"] is False + assert middle_info["runtime_cache_hit"] is False + assert revisit_info["runtime_cache_hit"] is True + assert runtime.cache_info() == { + "size": 2, + "hits": 1, + "misses": 2, + "max_cached_shapes": None, + } + finally: + runtime.clear() + + +@pytest.mark.gpu +def test_internal_fused_norms_and_explicit_precomputed_norms_are_correct(): + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("CUDA GPU required for exported-kernel correctness") + from flashlib_cake_kmeans import init + + rows = {row["label"]: row for row in BENCHMARK.FLASH_KMEANS_SHAPES} + row = rows["boundary_b1_n128_k256_d128"] + x, centroids = BENCHMARK._make_inputs(row, variant=250) + x_sq = (x.float() ** 2).sum(-1).contiguous() + c_sq = (centroids.float() ** 2).sum(-1).contiguous() + internal_out = torch.empty((1, 128), dtype=torch.int32, device=x.device) + explicit_out = torch.empty_like(internal_out) + x_explicit_out = torch.empty_like(internal_out) + c_explicit_out = torch.empty_like(internal_out) + runtime = init() + try: + internal_result, internal_info = runtime.compute( + x, + centroids, + out=internal_out, + return_info=True, + ) + explicit_result, explicit_info = runtime.compute( + x, + centroids, + out=explicit_out, + x_sq=x_sq, + c_sq=c_sq, + return_info=True, + ) + x_explicit_result, x_explicit_info = runtime.compute( + x, + centroids, + out=x_explicit_out, + x_sq=x_sq, + return_info=True, + ) + c_explicit_result, c_explicit_info = runtime.compute( + x, + centroids, + out=c_explicit_out, + c_sq=c_sq, + return_info=True, + ) + reference = BENCHMARK._reference_assign(x, centroids, chunk_rows=128) + torch.cuda.synchronize() + for result in ( + internal_result, + explicit_result, + x_explicit_result, + c_explicit_result, + ): + correctness = BENCHMARK._assignment_correctness( + result, + reference, + x, + centroids, + ) + assert correctness["correct"], correctness + assert internal_info["norm_mode"] == "fused_bf16_pair_row_norm:x_sq,c_sq" + assert internal_info["norm_compute_fields"] == ["x_sq", "c_sq"] + assert explicit_info["norm_mode"] == "explicit_precomputed" + assert explicit_info["norm_compute_fields"] == [] + assert x_explicit_info["norm_mode"] == "fused_bf16_pair_row_norm:c_sq" + assert x_explicit_info["norm_compute_fields"] == ["c_sq"] + assert c_explicit_info["norm_mode"] == "fused_bf16_pair_row_norm:x_sq" + assert c_explicit_info["norm_compute_fields"] == ["x_sq"] + assert explicit_info["runtime_cache_hit"] is False + assert x_explicit_info["runtime_cache_hit"] is False + assert c_explicit_info["runtime_cache_hit"] is False + assert runtime.cache_info()["size"] == 4 + finally: + runtime.clear() + + +@pytest.mark.gpu +def test_kmeans_runtime_isolates_tma_and_workspace_across_streams(): + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("CUDA GPU required for exported-kernel correctness") + from flashlib_cake_kmeans import init + + rows = {row["label"]: row for row in BENCHMARK.FLASH_KMEANS_SHAPES} + row = rows["boundary_b1_n128_k256_d128"] + stream_a = torch.cuda.Stream() + stream_b = torch.cuda.Stream() + runtime = init() + + def inputs(variant: int): + x, centroids = BENCHMARK._make_inputs(row, variant=variant) + return ( + x, + centroids, + (x.float() ** 2).sum(-1).contiguous(), + (centroids.float() ** 2).sum(-1).contiguous(), + torch.empty((int(row["B"]), int(row["N"])), dtype=torch.int32, device=x.device), + ) + + values_a = inputs(301) + values_b = inputs(302) + try: + for stream, values in ((stream_a, values_a), (stream_b, values_b)): + x, centroids, x_sq, c_sq, out = values + with torch.cuda.stream(stream): + runtime.compute( + x, + centroids, + x_sq=x_sq, + c_sq=c_sq, + out=out, + stream=stream, + ) + stream_a.synchronize() + stream_b.synchronize() + + outputs = [] + for stream, values in ((stream_a, values_a), (stream_b, values_b)): + x, centroids, x_sq, c_sq, out = values + with torch.cuda.stream(stream): + outputs.append( + runtime.compute( + x, + centroids, + x_sq=x_sq, + c_sq=c_sq, + out=out, + stream=stream, + ) + ) + stream_a.synchronize() + stream_b.synchronize() + for output, values in zip(outputs, (values_a, values_b), strict=True): + x, centroids, *_ = values + reference = BENCHMARK._reference_assign(x, centroids, chunk_rows=128) + correctness = BENCHMARK._assignment_correctness(output, reference, x, centroids) + assert correctness["correct"], correctness + assert runtime.cache_info()["size"] == 2 + finally: + runtime.clear() diff --git a/cake_exports/kmeans/tests/test_exported_kernels.py b/cake_exports/kmeans/tests/test_exported_kernels.py new file mode 100644 index 00000000..093dc9cb --- /dev/null +++ b/cake_exports/kmeans/tests/test_exported_kernels.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +import importlib +import json +import sys +import types +from pathlib import Path + +import pytest + + +ROOT = Path(__file__).resolve().parents[1] +SRC = ROOT / "src" +if str(SRC) not in sys.path: + sys.path.insert(0, str(SRC)) + +PACKAGE_NAME = 'flashlib_cake_kmeans' + + +def _manifest() -> dict: + return json.loads((SRC / PACKAGE_NAME / "manifest.json").read_text(encoding="utf-8")) + + +def test_manifest_matches_package_exports(): + pkg = importlib.import_module(PACKAGE_NAME) + manifest = _manifest() + names = [entry["name"] for entry in manifest["kernels"]] + + assert names + assert set(pkg.KERNELS) == set(names) + assert manifest["package"] == PACKAGE_NAME + package_exports = manifest.get("export_plan", {}).get("package_exports", {}) + for public_name in package_exports: + assert hasattr(pkg, public_name), public_name + assert public_name in pkg.__all__, public_name + + if "export_plan" in manifest: + entrypoints = manifest["export_plan"].get("entrypoints", {}) + assert set(entrypoints) == { + "python_interface", + "correctness_test", + "performance_benchmark", + } + for path in entrypoints.values(): + assert (ROOT / path).is_file(), path + + +def test_manifest_sources_exist_and_contain_symbols(): + manifest = _manifest() + package_dir = SRC / PACKAGE_NAME + + for entry in manifest["kernels"]: + source_path = package_dir / entry["source"] + assert source_path.is_file(), entry["source"] + source = source_path.read_text(encoding="utf-8") + assert entry["symbol"] in source + assert entry["parameters"] + + +def test_package_import_and_source_text_do_not_require_cuda_runtime(): + pkg = importlib.import_module(PACKAGE_NAME) + + for name, kernel in pkg.KERNELS.items(): + assert kernel.source_text().startswith("typedef "), name + assert kernel.parameters + + +def test_exported_repo_docs_and_benchmarks_exist(): + assert (ROOT / "README.md").is_file() + assert (ROOT / "RESULTS.md").is_file() + assert (ROOT / "benchmarks" / "benchmark_exported_kernels.py").is_file() + assert (ROOT / "benchmarks" / "benchmark_shapes.py").is_file() + assert (ROOT / "benchmarks" / "workload.py").is_file() + assert (SRC / PACKAGE_NAME / "tvm_ffi.py").is_file() + pyproject = (ROOT / "pyproject.toml").read_text(encoding="utf-8") + assert 'cake-std==0.1.13.dev20260704+g7b8dbc8' in pyproject + + +def test_tvm_ffi_adapter_registers_low_level_functions_without_import_time_dependency(monkeypatch): + pkg = importlib.import_module(PACKAGE_NAME) + registrations = {} + + class FakeTensor: + pass + + def register_global_func(name, function, *, override=False): + assert override is False + registrations[name] = function + + fake_tvm_ffi = types.SimpleNamespace( + Tensor=FakeTensor, + register_global_func=register_global_func, + get_raw_stream=lambda device: 0, + ) + monkeypatch.setitem(sys.modules, "tvm_ffi", fake_tvm_ffi) + + expected = pkg.tvm_ffi_function_names("export_test") + registered = pkg.register_tvm_ffi("export_test") + assert registered == expected + assert set(registrations) == set(expected) + assert registered == pkg.register_tvm_ffi("export_test") + assert len(registrations) == len(expected) + + +def test_benchmark_runtime_requires_cupti_without_event_or_wall_clock_fallback(): + source = (SRC / PACKAGE_NAME / "_benchmark.py").read_text(encoding="utf-8") + assert "activity_register_callbacks" in source + assert "torch.cuda.Event" not in source + assert "perf_counter" not in source + assert "active_union_times_ms" in source + assert "activity_counts" in source + assert "launch_activity_counts" in source + assert "kernel_activity_counts" in source + assert "submission_times_ms" in source + assert "synchronized_e2e_times_ms" in source + assert "cold_first_call_host_enqueue_ms" in source + assert "cold_first_call_synchronized_e2e_ms" in source + helper = source.split("def _complete_l2_flush_before_bracket", 1)[1].split( + "def _correlate", 1 + )[0] + assert helper.index("flusher.flush()") < helper.index("synchronize()") + measured_loop = source.split("for _ in range(bench_iters):", 1)[1] + assert measured_loop.index("_complete_l2_flush_before_bracket") < measured_loop.index( + "start = cupti.get_timestamp()" + ) + + +def test_benchmark_runtime_preserves_exact_activity_diagnostics(): + benchmark = importlib.import_module(f"{PACKAGE_NAME}._benchmark") + timing = benchmark._correlate( + [(0, 100, 10_000)], + [(10, 11, 1), (20, 21, 2)], + [(1_000, 4_000, 1), (2_000, 5_000, 2)], + ) + assert timing.gpu_span_ms == [0.004] + assert timing.kernel_sum_ms == [0.006] + assert timing.active_union_ms == [0.004] + assert timing.inter_kernel_gap_ms == [0.0] + assert timing.activity_count == [2] + assert timing.launch_activity_count == [2] + assert timing.kernel_activity_count == [2] + with pytest.raises(ValueError, match="must be 'cupti'"): + benchmark.BenchResult(times_ms=[1.0], backend="cuda_event") + + +def test_runtime_has_content_cache_and_launch_context(): + source = (SRC / PACKAGE_NAME / "_runtime.py").read_text(encoding="utf-8") + assert "def launch_context" in source + assert "def load_cached_kernel" in source + assert "_CUBIN_CACHE" in source + assert "_MODULE_CACHE" in source + assert "_KERNEL_CACHE" in source + kernels_source = (SRC / PACKAGE_NAME / "kernels.py").read_text(encoding="utf-8") + assert "self._arg_types = tuple" in kernels_source + assert "self._default_block" in kernels_source + + +def test_launch_argument_count_is_checked_before_compilation(): + pkg = importlib.import_module(PACKAGE_NAME) + kernel = next(iter(pkg.KERNELS.values())) + bad_args = [object()] * (len(kernel.parameters) + 1) + + with pytest.raises(TypeError, match="expects"): + kernel.launch(*bad_args, grid=(1, 1, 1)) +